add rate limit

app.py

@@ -6,19 +6,18 @@ import urllib.request
 from utils import model_name_mapping, urial_template, openai_base_request, DEFAULT_API_KEY
 from constant import js_code_label, HEADER_MD
 from openai import OpenAI
-
+import datetime
 # add logging info to console 
 logging.basicConfig(level=logging.INFO)
 
-
-
 URIAL_VERSION = "inst_1k_v4.help"
-
 URIAL_URL = f"https://raw.githubusercontent.com/Re-Align/URIAL/main/urial_prompts/{URIAL_VERSION}.txt"
 urial_prompt = urllib.request.urlopen(URIAL_URL).read().decode('utf-8')
 urial_prompt = urial_prompt.replace("```", '"""') # new version of URIAL uses """ instead of ```
 STOP_STRS = ['"""', '# Query:', '# Answer:']
 
+addr_limit_counter = {}
+LAST_UPDATE_TIME = datetime.datetime.now() 
 
 def respond(
     message,
@@ -28,9 +27,10 @@ def respond(
     top_p,
     rp,
     model_name,
-    together_api_key
+    together_api_key,
+    request:gr.Request
 ):  
-    global STOP_STRS, urial_prompt
+    global STOP_STRS, urial_prompt, LAST_UPDATE_TIME, addr_limit_counter
     rp = 1.0
     prompt = urial_template(urial_prompt, history, message)
     
@@ -42,14 +42,29 @@ def respond(
     else:
         api_key = DEFAULT_API_KEY
 
-    request = openai_base_request(prompt=prompt, model=_model_name, 
+    # headers = request.headers
+    # if already 24 hours passed, reset the counter
+    if datetime.datetime.now() - LAST_UPDATE_TIME > datetime.timedelta(days=1):
+        addr_limit_counter = {}
+        LAST_UPDATE_TIME = datetime.datetime.now()
+    host_addr = request.client.host
+    if host_addr not in addr_limit_counter:
+        addr_limit_counter[host_addr] = 0
+    if addr_limit_counter[host_addr] > 100:
+        return "You have reached the limit of 100 requests for today. Please use your own API key."
+
+    infer_request = openai_base_request(prompt=prompt, model=_model_name, 
                                    temperature=temperature, 
                                    max_tokens=max_tokens, 
                                    top_p=top_p, 
                                    repetition_penalty=rp,
                                    stop=STOP_STRS, api_key=api_key)  
+    addr_limit_counter[host_addr] += 1
+    logging.info(f"Requesting chat completion from OpenAI API with model {_model_name}")
+    logging.info(f"addr_limit_counter: {addr_limit_counter}; Last update time: {LAST_UPDATE_TIME};")
+
     response = ""
-    for msg in request:
+    for msg in infer_request:
         # print(msg.choices[0].delta.keys())
         token = msg.choices[0].delta["content"]
         should_stop = False
@@ -91,5 +106,5 @@ with gr.Blocks(gr.themes.Soft(), js=js_code_label) as demo:
     chat.chatbot.height = 550
     chat.chatbot.show_copy_button = True 
 
-if __name__ == "__main__":
+if __name__ == "__main__": 
     demo.launch(show_api=False)