mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-10-09 20:55:06 +02:00
Add infinite retry for starting up Slack bot (#540)
This commit is contained in:
@@ -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)
|
||||||
|
@@ -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(
|
||||||
|
@@ -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")
|
||||||
|
Reference in New Issue
Block a user