Push multi tenancy for slackbot (#2828)

* push multi tenancy for slackbot * move to utils * k * k --------- Co-authored-by: hagen-danswer <hagen@danswer.ai>
onyx-dot-app · Oct 17, 2024 · b169f78 · b169f78
1 parent e48086b
commit b169f78
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 90 deletions.
diff --git a/backend/danswer/danswerbot/slack/handlers/handle_buttons.py b/backend/danswer/danswerbot/slack/handlers/handle_buttons.py
@@ -4,9 +4,7 @@
 from slack_sdk import WebClient
 from slack_sdk.models.blocks import SectionBlock
 from slack_sdk.models.views import View
-from slack_sdk.socket_mode import SocketModeClient
 from slack_sdk.socket_mode.request import SocketModeRequest
-from sqlalchemy.orm import Session
 
 from danswer.configs.constants import MessageType
 from danswer.configs.constants import SearchFeedbackType
@@ -35,20 +33,22 @@
 from danswer.danswerbot.slack.utils import get_feedback_visibility
 from danswer.danswerbot.slack.utils import read_slack_thread
 from danswer.danswerbot.slack.utils import respond_in_thread
+from danswer.danswerbot.slack.utils import TenantSocketModeClient
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_sqlalchemy_engine
+from danswer.db.engine import get_session_with_tenant
 from danswer.db.feedback import create_chat_message_feedback
 from danswer.db.feedback import create_doc_retrieval_feedback
 from danswer.document_index.document_index_utils import get_both_index_names
 from danswer.document_index.factory import get_default_document_index
 from danswer.utils.logger import setup_logger
 
+
 logger = setup_logger()
 
 
 def handle_doc_feedback_button(
     req: SocketModeRequest,
-    client: SocketModeClient,
+    client: TenantSocketModeClient,
 ) -> None:
     if not (actions := req.payload.get("actions")):
         logger.error("Missing actions. Unable to build the source feedback view")
@@ -81,7 +81,7 @@ def handle_doc_feedback_button(
 
 def handle_generate_answer_button(
     req: SocketModeRequest,
-    client: SocketModeClient,
+    client: TenantSocketModeClient,
 ) -> None:
     channel_id = req.payload["channel"]["id"]
     channel_name = req.payload["channel"]["name"]
@@ -116,7 +116,7 @@ def handle_generate_answer_button(
         thread_ts=thread_ts,
     )
 
-    with Session(get_sqlalchemy_engine()) as db_session:
+    with get_session_with_tenant(client.tenant_id) as db_session:
         slack_bot_config = get_slack_bot_config_for_channel(
             channel_name=channel_name, db_session=db_session
         )
@@ -136,6 +136,7 @@ def handle_generate_answer_button(
             slack_bot_config=slack_bot_config,
             receiver_ids=None,
             client=client.web_client,
+            tenant_id=client.tenant_id,
             channel=channel_id,
             logger=logger,
             feedback_reminder_id=None,
@@ -150,12 +151,11 @@ def handle_slack_feedback(
     user_id_to_post_confirmation: str,
     channel_id_to_post_confirmation: str,
     thread_ts_to_post_confirmation: str,
+    tenant_id: str | None,
 ) -> None:
-    engine = get_sqlalchemy_engine()
-
     message_id, doc_id, doc_rank = decompose_action_id(feedback_id)
 
-    with Session(engine) as db_session:
+    with get_session_with_tenant(tenant_id) as db_session:
         if feedback_type in [LIKE_BLOCK_ACTION_ID, DISLIKE_BLOCK_ACTION_ID]:
             create_chat_message_feedback(
                 is_positive=feedback_type == LIKE_BLOCK_ACTION_ID,
@@ -232,7 +232,7 @@ def handle_slack_feedback(
 
 def handle_followup_button(
     req: SocketModeRequest,
-    client: SocketModeClient,
+    client: TenantSocketModeClient,
 ) -> None:
     action_id = None
     if actions := req.payload.get("actions"):
@@ -252,7 +252,7 @@ def handle_followup_button(
 
     tag_ids: list[str] = []
     group_ids: list[str] = []
-    with Session(get_sqlalchemy_engine()) as db_session:
+    with get_session_with_tenant(client.tenant_id) as db_session:
         channel_name, is_dm = get_channel_name_from_id(
             client=client.web_client, channel_id=channel_id
         )
@@ -295,7 +295,7 @@ def handle_followup_button(
 
 def get_clicker_name(
     req: SocketModeRequest,
-    client: SocketModeClient,
+    client: TenantSocketModeClient,
 ) -> str:
     clicker_name = req.payload.get("user", {}).get("name", "Someone")
     clicker_real_name = None
@@ -316,7 +316,7 @@ def get_clicker_name(
 
 def handle_followup_resolved_button(
     req: SocketModeRequest,
-    client: SocketModeClient,
+    client: TenantSocketModeClient,
     immediate: bool = False,
 ) -> None:
     channel_id = req.payload["container"]["channel_id"]

diff --git a/backend/danswer/danswerbot/slack/handlers/handle_message.py b/backend/danswer/danswerbot/slack/handlers/handle_message.py
@@ -2,7 +2,6 @@
 
 from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
-from sqlalchemy.orm import Session
 
 from danswer.configs.danswerbot_configs import DANSWER_BOT_FEEDBACK_REMINDER
 from danswer.configs.danswerbot_configs import DANSWER_REACT_EMOJI
@@ -19,7 +18,7 @@
 from danswer.danswerbot.slack.utils import respond_in_thread
 from danswer.danswerbot.slack.utils import slack_usage_report
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_sqlalchemy_engine
+from danswer.db.engine import get_session_with_tenant
 from danswer.db.models import SlackBotConfig
 from danswer.db.users import add_non_web_user_if_not_exists
 from danswer.utils.logger import setup_logger
@@ -110,6 +109,7 @@ def handle_message(
     slack_bot_config: SlackBotConfig | None,
     client: WebClient,
     feedback_reminder_id: str | None,
+    tenant_id: str | None,
 ) -> bool:
     """Potentially respond to the user message depending on filters and if an answer was generated
 
@@ -135,7 +135,9 @@ def handle_message(
         action = "slack_tag_message"
     elif is_bot_dm:
         action = "slack_dm_message"
-    slack_usage_report(action=action, sender_id=sender_id, client=client)
+    slack_usage_report(
+        action=action, sender_id=sender_id, client=client, tenant_id=tenant_id
+    )
 
     document_set_names: list[str] | None = None
     persona = slack_bot_config.persona if slack_bot_config else None
@@ -209,7 +211,7 @@ def handle_message(
     except SlackApiError as e:
         logger.error(f"Was not able to react to user message due to: {e}")
 
-    with Session(get_sqlalchemy_engine()) as db_session:
+    with get_session_with_tenant(tenant_id) as db_session:
         if message_info.email:
             add_non_web_user_if_not_exists(db_session, message_info.email)
 
@@ -235,5 +237,6 @@ def handle_message(
             channel=channel,
             logger=logger,
             feedback_reminder_id=feedback_reminder_id,
+            tenant_id=tenant_id,
         )
         return issue_with_regular_answer
diff --git a/backend/danswer/danswerbot/slack/handlers/handle_regular_answer.py b/backend/danswer/danswerbot/slack/handlers/handle_regular_answer.py
@@ -9,7 +9,6 @@
 from slack_sdk import WebClient
 from slack_sdk.models.blocks import DividerBlock
 from slack_sdk.models.blocks import SectionBlock
-from sqlalchemy.orm import Session
 
 from danswer.configs.app_configs import DISABLE_GENERATIVE_AI
 from danswer.configs.danswerbot_configs import DANSWER_BOT_ANSWER_GENERATION_TIMEOUT
@@ -33,7 +32,7 @@
 from danswer.danswerbot.slack.utils import respond_in_thread
 from danswer.danswerbot.slack.utils import SlackRateLimiter
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_sqlalchemy_engine
+from danswer.db.engine import get_session_with_tenant
 from danswer.db.models import Persona
 from danswer.db.models import SlackBotConfig
 from danswer.db.models import SlackBotResponseType
@@ -88,6 +87,7 @@ def handle_regular_answer(
     channel: str,
     logger: DanswerLoggingAdapter,
     feedback_reminder_id: str | None,
+    tenant_id: str | None,
     num_retries: int = DANSWER_BOT_NUM_RETRIES,
     answer_generation_timeout: int = DANSWER_BOT_ANSWER_GENERATION_TIMEOUT,
     thread_context_percent: float = DANSWER_BOT_TARGET_CHUNK_PERCENTAGE,
@@ -104,8 +104,7 @@ def handle_regular_answer(
     user = None
     if message_info.is_bot_dm:
         if message_info.email:
-            engine = get_sqlalchemy_engine()
-            with Session(engine) as db_session:
+            with get_session_with_tenant(tenant_id) as db_session:
                 user = get_user_by_email(message_info.email, db_session)
 
     document_set_names: list[str] | None = None
@@ -152,7 +151,7 @@ def _get_answer(new_message_request: DirectQARequest) -> OneShotQAResponse | Non
         max_document_tokens: int | None = None
         max_history_tokens: int | None = None
 
-        with Session(get_sqlalchemy_engine()) as db_session:
+        with get_session_with_tenant(tenant_id) as db_session:
             if len(new_message_request.messages) > 1:
                 if new_message_request.persona_config:
                     raise RuntimeError("Slack bot does not support persona config")
@@ -246,7 +245,7 @@ def _get_answer(new_message_request: DirectQARequest) -> OneShotQAResponse | Non
         )
 
         # Always apply reranking settings if it exists, this is the non-streaming flow
-        with Session(get_sqlalchemy_engine()) as db_session:
+        with get_session_with_tenant(tenant_id) as db_session:
             saved_search_settings = get_current_search_settings(db_session)
 
         # This includes throwing out answer via reflexion