Add infinite retry for starting up Slack bot (#540)

This commit is contained in:
Chris Weaver
2023-10-08 17:02:40 -07:00
committed by GitHub
parent ca74884bd7
commit f045bbed70
3 changed files with 28 additions and 19 deletions

View File

@@ -1,5 +1,6 @@
import logging import logging
import re import re
import time
from collections.abc import MutableMapping from collections.abc import MutableMapping
from typing import Any from typing import Any
from typing import cast from typing import cast
@@ -11,9 +12,9 @@ from slack_sdk.socket_mode.response import SocketModeResponse
from danswer.bots.slack.handlers.handle_feedback import handle_slack_feedback from danswer.bots.slack.handlers.handle_feedback import handle_slack_feedback
from danswer.bots.slack.handlers.handle_message import handle_message from danswer.bots.slack.handlers.handle_message import handle_message
from danswer.bots.slack.tokens import fetch_tokens
from danswer.bots.slack.utils import decompose_block_id from danswer.bots.slack.utils import decompose_block_id
from danswer.dynamic_configs.interface import ConfigNotFoundError from danswer.dynamic_configs.interface import ConfigNotFoundError
from danswer.server.slack_bot_management import get_tokens
from danswer.utils.logger import setup_logger from danswer.utils.logger import setup_logger
logger = setup_logger() logger = setup_logger()
@@ -22,6 +23,10 @@ logger = setup_logger()
_CHANNEL_ID = "channel_id" _CHANNEL_ID = "channel_id"
class MissingTokensException(Exception):
pass
class _ChannelIdAdapter(logging.LoggerAdapter): class _ChannelIdAdapter(logging.LoggerAdapter):
"""This is used to add the channel ID to all log messages """This is used to add the channel ID to all log messages
emitted in this file""" emitted in this file"""
@@ -40,9 +45,9 @@ def _get_socket_client() -> SocketModeClient:
# For more info on how to set this up, checkout the docs: # For more info on how to set this up, checkout the docs:
# https://docs.danswer.dev/slack_bot_setup # https://docs.danswer.dev/slack_bot_setup
try: try:
slack_bot_tokens = get_tokens() slack_bot_tokens = fetch_tokens()
except ConfigNotFoundError: except ConfigNotFoundError:
raise RuntimeError("Slack tokens not found") raise MissingTokensException("Slack tokens not found")
return SocketModeClient( return SocketModeClient(
# This app-level token will be used only for establishing a connection # This app-level token will be used only for establishing a connection
app_token=slack_bot_tokens.app_token, app_token=slack_bot_tokens.app_token,
@@ -222,14 +227,21 @@ def process_slack_event(client: SocketModeClient, req: SocketModeRequest) -> Non
# NOTE: we are using Web Sockets so that you can run this from within a firewalled VPC # NOTE: we are using Web Sockets so that you can run this from within a firewalled VPC
# without issue. # without issue.
if __name__ == "__main__": if __name__ == "__main__":
socket_client = _get_socket_client() try:
socket_client.socket_mode_request_listeners.append(process_slack_event) # type: ignore socket_client = _get_socket_client()
socket_client.socket_mode_request_listeners.append(process_slack_event) # type: ignore
# Establish a WebSocket connection to the Socket Mode servers # Establish a WebSocket connection to the Socket Mode servers
logger.info("Listening for messages from Slack...") logger.info("Listening for messages from Slack...")
socket_client.connect() socket_client.connect()
# Just not to stop this process # Just not to stop this process
from threading import Event from threading import Event
Event().wait() Event().wait()
except MissingTokensException:
# try again every 30 seconds. This is needed since the user may add tokens
# via the UI at any point in the programs lifecycle - if we just allow it to
# fail, then the user will need to restart the containers after adding tokens
logger.debug("Missing Slack Bot tokens - waiting 60 seconds and trying again")
time.sleep(60)

View File

@@ -2,7 +2,6 @@ import os
from typing import cast from typing import cast
from danswer.dynamic_configs import get_dynamic_config_store from danswer.dynamic_configs import get_dynamic_config_store
from danswer.dynamic_configs.interface import ConfigNotFoundError
from danswer.server.models import SlackBotTokens from danswer.server.models import SlackBotTokens
@@ -17,12 +16,9 @@ def fetch_tokens() -> SlackBotTokens:
return SlackBotTokens(app_token=app_token, bot_token=bot_token) return SlackBotTokens(app_token=app_token, bot_token=bot_token)
dynamic_config_store = get_dynamic_config_store() dynamic_config_store = get_dynamic_config_store()
try: return SlackBotTokens(
return SlackBotTokens( **cast(dict, dynamic_config_store.load(key=_SLACK_BOT_TOKENS_CONFIG_KEY))
**cast(dict, dynamic_config_store.load(key=_SLACK_BOT_TOKENS_CONFIG_KEY)) )
)
except ConfigNotFoundError as e:
raise ValueError from e
def save_tokens( def save_tokens(

View File

@@ -14,6 +14,7 @@ from danswer.db.slack_bot_config import fetch_slack_bot_configs
from danswer.db.slack_bot_config import insert_slack_bot_config from danswer.db.slack_bot_config import insert_slack_bot_config
from danswer.db.slack_bot_config import remove_slack_bot_config from danswer.db.slack_bot_config import remove_slack_bot_config
from danswer.db.slack_bot_config import update_slack_bot_config from danswer.db.slack_bot_config import update_slack_bot_config
from danswer.dynamic_configs.interface import ConfigNotFoundError
from danswer.server.models import DocumentSet from danswer.server.models import DocumentSet
from danswer.server.models import SlackBotConfig from danswer.server.models import SlackBotConfig
from danswer.server.models import SlackBotConfigCreationRequest from danswer.server.models import SlackBotConfigCreationRequest
@@ -169,5 +170,5 @@ def put_tokens(tokens: SlackBotTokens) -> None:
def get_tokens() -> SlackBotTokens: def get_tokens() -> SlackBotTokens:
try: try:
return fetch_tokens() return fetch_tokens()
except ValueError: except ConfigNotFoundError:
raise HTTPException(status_code=404, detail="No tokens found") raise HTTPException(status_code=404, detail="No tokens found")