From c20de7815f16c0a49605e00e8955e83bbc69b676 Mon Sep 17 00:00:00 2001
From: Federico Quattrocchio <federicoquattrocchio@gmail.com>
Date: Wed, 19 Jul 2023 10:59:55 -0300
Subject: [PATCH] Add removed code (#36)

---
 src/services/api_service.py    | 38 ++++++++++++++++++++++++++++------
 src/services/openai_service.py |  2 --
 src/services/slack_service.py  | 35 ++++++++++++++++++++++++-------
 3 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/src/services/api_service.py b/src/services/api_service.py
index 353525d..87019f3 100644
--- a/src/services/api_service.py
+++ b/src/services/api_service.py
@@ -1,18 +1,25 @@
+import json
 import os
-
+from urllib.request import Request, urlopen
 import requests
 
 BASE_URL = os.environ["API_BASE_URL"]
+SHARED_SECRET = os.environ["API_SHARED_SECRET"]
 
 
-def get_key(team_id):
-    url = f"{BASE_URL}/api/organization/get_key/{team_id}"
-    headers = {"X-Shared-Secret": os.environ["API_SHARED_SECRET"]}
+def get_team_data(team_id):
+    url = f"{BASE_URL}/api/organization/get_team_data/{team_id}"
+    headers = {"X-Shared-Secret": SHARED_SECRET}
     response = requests.get(url=url, headers=headers, timeout=30)
     data = response.json()
     if (data.get("error") is not None):
         raise Exception(data["error"])
-    return {"openai_key": data["openai_key"], "slack_bot_token": data["slack_bot_token"]}
+    return {
+        "openai_key": data["openai_key"],
+        "slack_bot_token": data["slack_bot_token"],
+        "has_reached_request_limit": data["has_reached_request_limit"],
+        "owner_email": data["owner_email"],
+    }
 
 
 def revoke_token(team_id):
@@ -31,4 +38,23 @@ def revoke_token(team_id):
     data = response.json()
     if (data.get("error") is not None):
         raise Exception(data["error"])
-    return
\ No newline at end of file
+    return
+
+def increment_request_count(team_id):
+    url = f"{BASE_URL}/api/slack/increment_request_count"
+    headers = {
+        "X-Shared-Secret": SHARED_SECRET,
+        "Content-Type": "application/json",
+    }
+    response = requests.post(
+        url=url,
+        headers=headers,
+        json={
+            "team_id": team_id,
+        },
+        timeout=30
+    )
+    data = response.json()
+    if (data.get("error") is not None):
+        raise Exception(data["error"])
+    return
diff --git a/src/services/openai_service.py b/src/services/openai_service.py
index e0b8f23..eca3262 100644
--- a/src/services/openai_service.py
+++ b/src/services/openai_service.py
@@ -1,7 +1,5 @@
 import openai
 from openai.error import AuthenticationError, RateLimitError
-import os
-from services.api_service import get_key
 
 # openai.api_key = os.environ["OPENAI_API_KEY"]
 
diff --git a/src/services/slack_service.py b/src/services/slack_service.py
index 44ead5e..aaa771b 100644
--- a/src/services/slack_service.py
+++ b/src/services/slack_service.py
@@ -1,10 +1,11 @@
+from threading import Thread
 from slack_bolt import App
 import time
 import os
 import re
 from services.openai_service import respond_to_user
 from lib.retry import retry
-from .api_service import get_key, revoke_token
+from .api_service import get_team_data, increment_request_count, revoke_token
 
 # grabs the credentials from .env directly
 slack_app = App()
@@ -134,19 +135,35 @@ def process_event_payload(payload):
     team_id = event.get("team")
     user = event.get("user")
     try:
-        keys = get_key(team_id)
         thread_to_reply = thread_ts
         if thread_ts != ts:
             thread_to_reply = ts
 
+        # Get neccessary keys
+        team_data = get_team_data(team_id)
+        slack_bot_token = team_data["slack_bot_token"]
+        openAi_key = team_data["openai_key"] if team_data["openai_key"] else os.environ["OPENAI_API_KEY"]
+
+        # Check quota
+        if team_data["has_reached_request_limit"] == True:
+            send_message(
+                channel,
+                thread_to_reply,
+                f"It appears you've exceeded the usage limit. To continue enjoying our services without interruption, kindly get in touch with your organization's administrator on {team_data['owner_email']} and request for a subscription upgrade.",
+                slack_bot_token
+            )
+            return
+
+        # Send 'thinking' message while we process the request
         msg_ts = send_message(
             channel,
             thread_to_reply,
             "*Thinking...*",
-            keys["slack_bot_token"]
+            slack_bot_token
         )
 
-        username = get_user_name(user, keys["slack_bot_token"])
+        # Get messages in thread
+        username = get_user_name(user, slack_bot_token)
         messages = [{
             "role": "user",
             "content": text + ". " + username,
@@ -157,17 +174,19 @@ def process_event_payload(payload):
                 get_thread_messages_with_usernames_json(
                     channel,
                     thread_ts,
-                    keys["slack_bot_token"]
+                    slack_bot_token
                 )
                 or messages
             )
 
         start_time = time.perf_counter()
-        response = respond_to_user(messages, keys["openai_key"])
+        response = respond_to_user(messages, openAi_key)
         end_time = time.perf_counter()
-        print(f"response generated in {round(end_time - start_time, 2)}s")
 
-        return update_message(channel, thread_to_reply, msg_ts, response, keys["slack_bot_token"])
+        print(f"response generated in {round(end_time - start_time, 2)}s")
+        Thread(target=increment_request_count, args=(team_id,)).start()
+        
+        return update_message(channel, thread_to_reply, msg_ts, response, slack_bot_token)
     except Exception as error:
         # Improve error handling
         print(error)