added token limits and error handling for long prompts

src/perfguru.py

@@ -3,6 +3,7 @@ import json
 import os
 import datetime
 from itertools import zip_longest
+import tiktoken
 
 from models import select_random_model
 from rag import select_random_formatter
@@ -19,11 +20,23 @@ def code_upload(code_file_select):
         return gr.Button(interactive=True)
 
 
+def token_limit_getter(model: str) -> int:
+    with open("token_limits.json", "r") as f:
+        token_limits = json.load(f)
+    if model in token_limits:
+        return token_limits[model]
+    return int(2e6)
+
+
 def chat_with_llms(prompt, code_files, profile_file, profile_type):
     model1 = select_random_model()
     model2 = select_random_model()
     formatter1 = select_random_formatter()
     formatter2 = select_random_formatter()
+    encoder1 = None
+    encoder2 = None
+    num_tokens_1 = 0
+    num_tokens_2 = 0
 
     print(f"Selected models: {model1.name} and {model2.name}")
 
@@ -33,6 +46,22 @@ def chat_with_llms(prompt, code_files, profile_file, profile_type):
     if formatted1 is None or formatted2 is None:
         error_helper("Failed to format prompt. Please try again.")
     
+    if model1.name[:3] == "gpt":
+        encoder1 = tiktoken.encoding_for_model(model1.name)
+    if model2.name[:3] == "gpt":
+        encoder2 = tiktoken.encoding_for_model(model2.name)
+
+    if encoder1:
+        num_tokens_1 = encoder1.encode(formatted1)
+    if encoder2:
+        num_tokens_2 = encoder2.encode(formatted2)
+    
+    token_limit_1 = token_limit_getter(model1.name)
+    token_limit_2 = token_limit_getter(model2.name)
+
+    if num_tokens_1 >= token_limit_1 or num_tokens_2 >= token_limit_2:
+        error_helper("Prompt is too long. Please try again.")
+
     response1 = model1.get_response(formatted1)
     response2 = model2.get_response(formatted2)
 

src/token_limits.json

@@ -0,0 +1,9 @@
+{
+    "gpt-4o": 128000,
+    "gpt-4o-mini": 128000,
+    "gpt-4-turbo": 128000,
+    "gpt-4": 8192,
+    "gpt-3.5-turbo": 16385,
+    "gemini-1.5-flash": 1048576,
+    "gemini-1.5-pro": 2097152
+}