From 44905d36e5a6648a525ac9b6ab2f88ee0c0e3c60 Mon Sep 17 00:00:00 2001
From: Yuhong Sun <yuhongsun96@gmail.com>
Date: Mon, 15 Jan 2024 21:47:24 -0800
Subject: [PATCH] Slack Rate Limit Options (#952)

---
 backend/danswer/configs/danswerbot_configs.py    |  6 +++---
 backend/danswer/danswerbot/slack/utils.py        | 11 +++++++++--
 deployment/docker_compose/docker-compose.dev.yml |  2 ++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/backend/danswer/configs/danswerbot_configs.py b/backend/danswer/configs/danswerbot_configs.py
index 606dcb0cf..6fd610f99 100644
--- a/backend/danswer/configs/danswerbot_configs.py
+++ b/backend/danswer/configs/danswerbot_configs.py
@@ -51,7 +51,7 @@ ENABLE_DANSWERBOT_REFLEXION = (
 # Currently not support chain of thought, probably will add back later
 DANSWER_BOT_DISABLE_COT = True
 
-# Maximum Questions Per Minute
-DANSWER_BOT_MAX_QPM = int(os.environ.get("DANSWER_BOT_MAX_QPM", "100"))
+# Maximum Questions Per Minute, Default Uncapped
+DANSWER_BOT_MAX_QPM = int(os.environ.get("DANSWER_BOT_MAX_QPM") or 0) or None
 # Maximum time to wait when a question is queued
-DANSWER_BOT_MAX_WAIT_TIME = int(os.environ.get("DANSWER_BOT_MAX_WAIT_TIME", "180"))
+DANSWER_BOT_MAX_WAIT_TIME = int(os.environ.get("DANSWER_BOT_MAX_WAIT_TIME") or 180)
diff --git a/backend/danswer/danswerbot/slack/utils.py b/backend/danswer/danswerbot/slack/utils.py
index b31c06b91..b54c14a42 100644
--- a/backend/danswer/danswerbot/slack/utils.py
+++ b/backend/danswer/danswerbot/slack/utils.py
@@ -355,7 +355,7 @@ def read_slack_thread(
 
 class SlackRateLimiter:
     def __init__(self) -> None:
-        self.max_qpm = DANSWER_BOT_MAX_QPM
+        self.max_qpm: int | None = DANSWER_BOT_MAX_QPM
         self.max_wait_time = DANSWER_BOT_MAX_WAIT_TIME
         self.active_question = 0
         self.last_reset_time = time.time()
@@ -374,11 +374,15 @@ class SlackRateLimiter:
             client=client,
             channel=channel,
             receiver_ids=None,
-            text=f"Your question has been queued. You are in position {position}... please wait a moment :loading:",
+            text=f"Your question has been queued. You are in position {position}.\n"
+            f"Please wait a moment :hourglass_flowing_sand:",
             thread_ts=thread_ts,
         )
 
     def is_available(self) -> bool:
+        if self.max_qpm is None:
+            return True
+
         self.refill()
         return self.active_question < self.max_qpm
 
@@ -393,6 +397,9 @@ class SlackRateLimiter:
         return func_randid, position
 
     def waiter(self, func_randid: int) -> None:
+        if self.max_qpm is None:
+            return
+
         wait_time = 0
         while (
             self.active_question >= self.max_qpm
diff --git a/deployment/docker_compose/docker-compose.dev.yml b/deployment/docker_compose/docker-compose.dev.yml
index 0553b7c17..56deb78d6 100644
--- a/deployment/docker_compose/docker-compose.dev.yml
+++ b/deployment/docker_compose/docker-compose.dev.yml
@@ -139,6 +139,8 @@ services:
       - DANSWER_BOT_RESPOND_EVERY_CHANNEL=${DANSWER_BOT_RESPOND_EVERY_CHANNEL:-}
       - DANSWER_BOT_DISABLE_COT=${DANSWER_BOT_DISABLE_COT:-}  # Currently unused
       - NOTIFY_SLACKBOT_NO_ANSWER=${NOTIFY_SLACKBOT_NO_ANSWER:-}
+      - DANSWER_BOT_MAX_QPM=${DANSWER_BOT_MAX_QPM:-}
+      - DANSWER_BOT_MAX_WAIT_TIME=${DANSWER_BOT_MAX_WAIT_TIME:-}
       # Logging
       # Leave this on pretty please? Nothing sensitive is collected!
       # https://docs.danswer.dev/more/telemetry