highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-06-27 08:21:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

adjust default postgres schema for slack listener (#3088)

Browse Source
This commit is contained in:
pablodanswer 2024-11-08 10:00:44 -08:00 committed by GitHub
parent f6d8f5ca89
commit cf5d394d39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 57 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

105
backend/danswer/danswerbot/slack/listener.py
View File

@ -77,6 +77,7 @@ from danswer.utils.logger import setup_logger
from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
from shared_configs.configs import MODEL_SERVER_HOST from shared_configs.configs import MODEL_SERVER_HOST
from shared_configs.configs import MODEL_SERVER_PORT from shared_configs.configs import MODEL_SERVER_PORT
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
from shared_configs.configs import SLACK_CHANNEL_ID from shared_configs.configs import SLACK_CHANNEL_ID
logger = setup_logger() logger = setup_logger()
@ -189,59 +190,67 @@ class SlackbotHandler:
continue continue
logger.debug(f"Acquired lock for tenant {tenant_id}") logger.debug(f"Acquired lock for tenant {tenant_id}")
with get_session_with_tenant(tenant_id) as db_session: token = CURRENT_TENANT_ID_CONTEXTVAR.set(
try: tenant_id or POSTGRES_DEFAULT_SCHEMA
logger.debug( )
f"Setting tenant ID context variable for tenant {tenant_id}" try:
) with get_session_with_tenant(tenant_id) as db_session:
token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id or "public") try:
slack_bot_tokens = fetch_tokens() logger.debug(
logger.debug(f"Fetched Slack bot tokens for tenant {tenant_id}") f"Setting tenant ID context variable for tenant {tenant_id}"
CURRENT_TENANT_ID_CONTEXTVAR.reset(token) )
logger.debug( slack_bot_tokens = fetch_tokens()
f"Reset tenant ID context variable for tenant {tenant_id}" logger.debug(f"Fetched Slack bot tokens for tenant {tenant_id}")
) logger.debug(
f"Reset tenant ID context variable for tenant {tenant_id}"
)
if not slack_bot_tokens: if not slack_bot_tokens:
logger.debug(f"No Slack bot token found for tenant {tenant_id}") logger.debug(
f"No Slack bot token found for tenant {tenant_id}"
)
if tenant_id in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id].close())
del self.socket_clients[tenant_id]
del self.slack_bot_tokens[tenant_id]
continue
if (
tenant_id not in self.slack_bot_tokens
or slack_bot_tokens != self.slack_bot_tokens[tenant_id]
):
if tenant_id in self.slack_bot_tokens:
logger.info(
f"Slack Bot tokens have changed for tenant {tenant_id} - reconnecting"
)
else:
search_settings = get_current_search_settings(
db_session
)
embedding_model = EmbeddingModel.from_db_model(
search_settings=search_settings,
server_host=MODEL_SERVER_HOST,
server_port=MODEL_SERVER_PORT,
)
warm_up_bi_encoder(embedding_model=embedding_model)
self.slack_bot_tokens[tenant_id] = slack_bot_tokens
if tenant_id in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id].close())
self.start_socket_client(tenant_id, slack_bot_tokens)
except KvKeyNotFoundError:
logger.debug(f"Missing Slack Bot tokens for tenant {tenant_id}")
if tenant_id in self.socket_clients: if tenant_id in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id].close()) asyncio.run(self.socket_clients[tenant_id].close())
del self.socket_clients[tenant_id] del self.socket_clients[tenant_id]
del self.slack_bot_tokens[tenant_id] del self.slack_bot_tokens[tenant_id]
continue except Exception as e:
logger.exception(f"Error handling tenant {tenant_id}: {e}")
if ( finally:
tenant_id not in self.slack_bot_tokens CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
or slack_bot_tokens != self.slack_bot_tokens[tenant_id]
):
if tenant_id in self.slack_bot_tokens:
logger.info(
f"Slack Bot tokens have changed for tenant {tenant_id} - reconnecting"
)
else:
search_settings = get_current_search_settings(db_session)
embedding_model = EmbeddingModel.from_db_model(
search_settings=search_settings,
server_host=MODEL_SERVER_HOST,
server_port=MODEL_SERVER_PORT,
)
warm_up_bi_encoder(embedding_model=embedding_model)
self.slack_bot_tokens[tenant_id] = slack_bot_tokens
if tenant_id in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id].close())
self.start_socket_client(tenant_id, slack_bot_tokens)
except KvKeyNotFoundError:
logger.debug(f"Missing Slack Bot tokens for tenant {tenant_id}")
if tenant_id in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id].close())
del self.socket_clients[tenant_id]
del self.slack_bot_tokens[tenant_id]
except Exception as e:
logger.exception(f"Error handling tenant {tenant_id}: {e}")
def send_heartbeats(self) -> None: def send_heartbeats(self) -> None:
current_time = int(time.time()) current_time = int(time.time())