Skip to content

Commit

Permalink
tested
Browse files Browse the repository at this point in the history
  • Loading branch information
hagen-danswer committed Nov 19, 2024
1 parent 4062069 commit 41ea202
Show file tree
Hide file tree
Showing 6 changed files with 70 additions and 66 deletions.
2 changes: 1 addition & 1 deletion backend/danswer/danswerbot/slack/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
]


def get_slack_channel_config_for_app_and_channel(
def get_slack_channel_config_for_bot_and_channel(
db_session: Session,
slack_bot_id: int,
channel_name: str | None,
Expand Down
10 changes: 5 additions & 5 deletions backend/danswer/danswerbot/slack/handlers/handle_buttons.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from danswer.connectors.slack.utils import make_slack_api_rate_limited
from danswer.danswerbot.slack.blocks import build_follow_up_resolved_blocks
from danswer.danswerbot.slack.blocks import get_document_feedback_blocks
from danswer.danswerbot.slack.config import get_slack_channel_config_for_app_and_channel
from danswer.danswerbot.slack.config import get_slack_channel_config_for_bot_and_channel
from danswer.danswerbot.slack.constants import DISLIKE_BLOCK_ACTION_ID
from danswer.danswerbot.slack.constants import FeedbackVisibility
from danswer.danswerbot.slack.constants import LIKE_BLOCK_ACTION_ID
Expand Down Expand Up @@ -117,9 +117,9 @@ def handle_generate_answer_button(
)

with get_session_with_tenant(client.tenant_id) as db_session:
slack_channel_config = get_slack_channel_config_for_app_and_channel(
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.app_id,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)

Expand Down Expand Up @@ -258,9 +258,9 @@ def handle_followup_button(
channel_name, is_dm = get_channel_name_from_id(
client=client.web_client, channel_id=channel_id
)
slack_channel_config = get_slack_channel_config_for_app_and_channel(
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.app_id,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)
if slack_channel_config:
Expand Down
90 changes: 48 additions & 42 deletions backend/danswer/danswerbot/slack/listener.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from danswer.configs.danswerbot_configs import DANSWER_BOT_RESPOND_EVERY_CHANNEL
from danswer.configs.danswerbot_configs import NOTIFY_SLACKBOT_NO_ANSWER
from danswer.connectors.slack.utils import expert_info_from_slack_id
from danswer.danswerbot.slack.config import get_slack_channel_config_for_app_and_channel
from danswer.danswerbot.slack.config import get_slack_channel_config_for_bot_and_channel
from danswer.danswerbot.slack.config import MAX_TENANTS_PER_POD
from danswer.danswerbot.slack.config import TENANT_ACQUISITION_INTERVAL
from danswer.danswerbot.slack.config import TENANT_HEARTBEAT_EXPIRATION
Expand Down Expand Up @@ -59,7 +59,7 @@
from danswer.danswerbot.slack.utils import check_message_limit
from danswer.danswerbot.slack.utils import decompose_action_id
from danswer.danswerbot.slack.utils import get_channel_name_from_id
from danswer.danswerbot.slack.utils import get_danswer_bot_app_id
from danswer.danswerbot.slack.utils import get_danswer_bot_slack_bot_id
from danswer.danswerbot.slack.utils import read_slack_thread
from danswer.danswerbot.slack.utils import remove_danswer_bot_tag
from danswer.danswerbot.slack.utils import rephrase_slack_message
Expand Down Expand Up @@ -117,7 +117,7 @@ class SlackbotHandler:
def __init__(self) -> None:
logger.info("Initializing SlackbotHandler")
self.tenant_ids: Set[str | None] = set()
# The keys for these dictionaries are tuples of (tenant_id, app_id)
# The keys for these dictionaries are tuples of (tenant_id, slack_bot_id)
self.socket_clients: Dict[tuple[str | None, int], TenantSocketModeClient] = {}
self.slack_bot_tokens: Dict[tuple[str | None, int], SlackBotTokens] = {}

Expand Down Expand Up @@ -176,38 +176,38 @@ def heartbeat_loop(self) -> None:
self._shutdown_event.wait(timeout=TENANT_HEARTBEAT_INTERVAL)

def _manage_clients_per_tenant(
self, db_session: Session, tenant_id: str | None, app: SlackBot
self, db_session: Session, tenant_id: str | None, bot: SlackBot
) -> None:
slack_bot_tokens = SlackBotTokens(
bot_token=app.bot_token,
app_token=app.app_token,
bot_token=bot.bot_token,
app_token=bot.app_token,
)

logger.debug(
f"Setting tenant ID context variable for tenant {tenant_id}, app {app.id}"
f"Setting tenant ID context variable for tenant {tenant_id}, bot {bot.id}"
)

# If the tokens are not set, we need to close the socket client and delete the tokens
# for the tenant and app
if not slack_bot_tokens:
logger.debug(
f"No Slack bot token found for tenant {tenant_id}, app {app.id}"
f"No Slack bot token found for tenant {tenant_id}, bot {bot.id}"
)
if (tenant_id, app.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, app.id].close())
del self.socket_clients[tenant_id, app.id]
del self.slack_bot_tokens[tenant_id, app.id]
if (tenant_id, bot.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, bot.id].close())
del self.socket_clients[tenant_id, bot.id]
del self.slack_bot_tokens[tenant_id, bot.id]
return

if (
tenant_id,
app.id,
bot.id,
) not in self.slack_bot_tokens or slack_bot_tokens != self.slack_bot_tokens[
(tenant_id, app.id)
(tenant_id, bot.id)
]:
if (tenant_id, app.id) in self.slack_bot_tokens:
if (tenant_id, bot.id) in self.slack_bot_tokens:
logger.info(
f"Slack Bot tokens have changed for tenant {tenant_id}, app {app.id} - reconnecting"
f"Slack Bot tokens have changed for tenant {tenant_id}, bot {bot.id} - reconnecting"
)
else:
search_settings = get_current_search_settings(db_session)
Expand All @@ -218,12 +218,12 @@ def _manage_clients_per_tenant(
)
warm_up_bi_encoder(embedding_model=embedding_model)

self.slack_bot_tokens[(tenant_id, app.id)] = slack_bot_tokens
self.slack_bot_tokens[(tenant_id, bot.id)] = slack_bot_tokens

if (tenant_id, app.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, app.id].close())
if (tenant_id, bot.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, bot.id].close())

self.start_socket_client(app.id, tenant_id, slack_bot_tokens)
self.start_socket_client(bot.id, tenant_id, slack_bot_tokens)

def acquire_tenants(self) -> None:
tenant_ids = get_all_tenant_ids()
Expand Down Expand Up @@ -269,20 +269,20 @@ def acquire_tenants(self) -> None:
try:
with get_session_with_tenant(tenant_id) as db_session:
try:
apps = fetch_slack_bots(db_session=db_session)
for app in apps:
bots = fetch_slack_bots(db_session=db_session)
for bot in bots:
self._manage_clients_per_tenant(
db_session=db_session,
tenant_id=tenant_id,
app=app,
bot=bot,
)

except KvKeyNotFoundError:
logger.debug(f"Missing Slack Bot tokens for tenant {tenant_id}")
if (tenant_id, app.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, app.id].close())
del self.socket_clients[tenant_id, app.id]
del self.slack_bot_tokens[tenant_id, app.id]
if (tenant_id, bot.id) in self.socket_clients:
asyncio.run(self.socket_clients[tenant_id, bot.id].close())
del self.socket_clients[tenant_id, bot.id]
del self.slack_bot_tokens[tenant_id, bot.id]
except Exception as e:
logger.exception(f"Error handling tenant {tenant_id}: {e}")
finally:
Expand All @@ -301,30 +301,36 @@ def send_heartbeats(self) -> None:
)

def start_socket_client(
self, app_id: int, tenant_id: str | None, slack_bot_tokens: SlackBotTokens
self, slack_bot_id: int, tenant_id: str | None, slack_bot_tokens: SlackBotTokens
) -> None:
logger.info(f"Starting socket client for tenant: {tenant_id}, app: {app_id}")
logger.info(
f"Starting socket client for tenant: {tenant_id}, app: {slack_bot_id}"
)
socket_client: TenantSocketModeClient = _get_socket_client(
slack_bot_tokens, tenant_id, app_id
slack_bot_tokens, tenant_id, slack_bot_id
)

# Append the event handler
process_slack_event = create_process_slack_event()
socket_client.socket_mode_request_listeners.append(process_slack_event) # type: ignore

# Establish a WebSocket connection to the Socket Mode servers
logger.info(f"Connecting socket client for tenant: {tenant_id}, app: {app_id}")
logger.info(
f"Connecting socket client for tenant: {tenant_id}, app: {slack_bot_id}"
)
socket_client.connect()
self.socket_clients[tenant_id, app_id] = socket_client
self.socket_clients[tenant_id, slack_bot_id] = socket_client
self.tenant_ids.add(tenant_id)
logger.info(f"Started SocketModeClient for tenant: {tenant_id}, app: {app_id}")
logger.info(
f"Started SocketModeClient for tenant: {tenant_id}, app: {slack_bot_id}"
)

def stop_socket_clients(self) -> None:
logger.info(f"Stopping {len(self.socket_clients)} socket clients")
for (tenant_id, app_id), client in self.socket_clients.items():
for (tenant_id, slack_bot_id), client in self.socket_clients.items():
asyncio.run(client.close())
logger.info(
f"Stopped SocketModeClient for tenant: {tenant_id}, app: {app_id}"
f"Stopped SocketModeClient for tenant: {tenant_id}, app: {slack_bot_id}"
)

def shutdown(self, signum: int | None, frame: FrameType | None) -> None:
Expand Down Expand Up @@ -409,7 +415,7 @@ def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -
)
return False

bot_tag_id = get_danswer_bot_app_id(client.web_client)
bot_tag_id = get_danswer_bot_slack_bot_id(client.web_client)
if event_type == "message":
is_dm = event.get("channel_type") == "im"
is_tagged = bot_tag_id and bot_tag_id in msg
Expand All @@ -432,9 +438,9 @@ def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -
)

with get_session_with_tenant(client.tenant_id) as db_session:
slack_channel_config = get_slack_channel_config_for_app_and_channel(
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.app_id,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)
# If DanswerBot is not specifically tagged and the channel is not set to respond to bots, ignore the message
Expand Down Expand Up @@ -645,9 +651,9 @@ def process_message(
token = CURRENT_TENANT_ID_CONTEXTVAR.set(client.tenant_id)
try:
with get_session_with_tenant(client.tenant_id) as db_session:
slack_channel_config = get_slack_channel_config_for_app_and_channel(
slack_channel_config = get_slack_channel_config_for_bot_and_channel(
db_session=db_session,
slack_bot_id=client.app_id,
slack_bot_id=client.slack_bot_id,
channel_name=channel_name,
)

Expand Down Expand Up @@ -753,7 +759,7 @@ def process_slack_event(


def _get_socket_client(
slack_bot_tokens: SlackBotTokens, tenant_id: str | None, app_id: int
slack_bot_tokens: SlackBotTokens, tenant_id: str | None, slack_bot_id: int
) -> TenantSocketModeClient:
# For more info on how to set this up, checkout the docs:
# https://docs.danswer.dev/slack_bot_setup
Expand All @@ -762,7 +768,7 @@ def _get_socket_client(
app_token=slack_bot_tokens.app_token,
web_client=WebClient(token=slack_bot_tokens.bot_token),
tenant_id=tenant_id,
app_id=app_id,
slack_bot_id=slack_bot_id,
)


Expand Down
24 changes: 13 additions & 11 deletions backend/danswer/danswerbot/slack/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,16 +46,16 @@
logger = setup_logger()


_DANSWER_BOT_APP_ID: str | None = None
_DANSWER_BOT_SLACK_BOT_ID: str | None = None
_DANSWER_BOT_MESSAGE_COUNT: int = 0
_DANSWER_BOT_COUNT_START_TIME: float = time.time()


def get_danswer_bot_app_id(web_client: WebClient) -> Any:
global _DANSWER_BOT_APP_ID
if _DANSWER_BOT_APP_ID is None:
_DANSWER_BOT_APP_ID = web_client.auth_test().get("user_id")
return _DANSWER_BOT_APP_ID
def get_danswer_bot_slack_bot_id(web_client: WebClient) -> Any:
global _DANSWER_BOT_SLACK_BOT_ID
if _DANSWER_BOT_SLACK_BOT_ID is None:
_DANSWER_BOT_SLACK_BOT_ID = web_client.auth_test().get("user_id")
return _DANSWER_BOT_SLACK_BOT_ID


def check_message_limit() -> bool:
Expand Down Expand Up @@ -136,7 +136,7 @@ def update_emote_react(


def remove_danswer_bot_tag(message_str: str, client: WebClient) -> str:
bot_tag_id = get_danswer_bot_app_id(web_client=client)
bot_tag_id = get_danswer_bot_slack_bot_id(web_client=client)
return re.sub(rf"<@{bot_tag_id}>\s", "", message_str)


Expand Down Expand Up @@ -431,9 +431,9 @@ def read_slack_thread(
)
message_type = MessageType.USER
else:
self_app_id = get_danswer_bot_app_id(client)
self_slack_bot_id = get_danswer_bot_slack_bot_id(client)

if reply.get("user") == self_app_id:
if reply.get("user") == self_slack_bot_id:
# DanswerBot response
message_type = MessageType.ASSISTANT
user_sem_id = "Assistant"
Expand Down Expand Up @@ -576,7 +576,9 @@ def get_feedback_visibility() -> FeedbackVisibility:


class TenantSocketModeClient(SocketModeClient):
def __init__(self, tenant_id: str | None, app_id: int, *args: Any, **kwargs: Any):
def __init__(
self, tenant_id: str | None, slack_bot_id: int, *args: Any, **kwargs: Any
):
super().__init__(*args, **kwargs)
self.tenant_id = tenant_id
self.app_id = app_id
self.slack_bot_id = slack_bot_id
8 changes: 2 additions & 6 deletions web/src/app/admin/bot/[id]/page.tsx
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { AdminPageTitle } from "@/components/admin/Title";
import { CPUIcon } from "@/components/icons/icons";
import { SourceIcon } from "@/components/SourceIcon";
import { SlackChannelConfigCreationForm } from "../SlackChannelConfigCreationForm";
import { fetchSS } from "@/lib/utilsSS";
import { ErrorCallout } from "@/components/ErrorCallout";
Expand Down Expand Up @@ -86,15 +87,10 @@ async function EditslackChannelConfigPage(props: {

<BackButton />
<AdminPageTitle
icon={<CPUIcon size={32} />}
icon={<SourceIcon sourceType={"slack"} iconSize={32} />}
title="Edit Slack Channel Config"
/>

<Text className="mb-8">
Edit the existing configuration below! This config will determine how
DanswerBot behaves in the specified channels.
</Text>

<SlackChannelConfigCreationForm
slack_bot_id={slackChannelConfig.slack_bot_id}
documentSets={documentSets}
Expand Down
2 changes: 1 addition & 1 deletion web/src/app/admin/bot/hooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ export const useSlackChannelConfigs = () => {
};

export const useSlackBots = () => {
const url = "/api/manage/admin/slack-app/apps";
const url = "/api/manage/admin/slack-app/bots";
const swrResponse = useSWR<SlackBot[]>(url, errorHandlingFetcher);

return {
Expand Down

0 comments on commit 41ea202

Please sign in to comment.