tokening

generationhelper.py

@@ -1,5 +1,7 @@
 import os
 from groq import Groq
+import time
+import tiktoken
 
 groq_token = os.getenv("GROQ_TOKEN")
 
@@ -7,6 +9,35 @@ groq_client = Groq(
     api_key = groq_token
 )
 
+# Initialize token counter and timestamp
+tokens_used = 0
+start_time = time.time()
+
+def Count_tokens(text: str, model="gpt-3.5-turbo"):
+    """Counts tokens in the given text using tiktoken."""
+    enc = tiktoken.encoding_for_model(model)
+    return len(enc.encode(text))
+
+def Enforce_token_limit(prompt, max_tokens_per_minute=6000):
+    """Ensures that token usage stays within the allowed rate limit."""
+    global tokens_used, start_time
+
+    tokens = Count_tokens(prompt)
+    elapsed_time = time.time() - start_time
+
+    # If the token limit is exceeded, wait until the reset
+    if tokens_used + tokens > max_tokens_per_minute:
+        if elapsed_time < 60:
+            sleep_time = 60 - elapsed_time
+            print(f"Rate limit reached! Sleeping for {sleep_time:.2f} seconds...")
+            time.sleep(sleep_time)
+
+        # Reset counter after sleeping
+        tokens_used = 0
+        start_time = time.time()
+
+    # Update token count
+    tokens_used += tokens
 
 
 def GenerateAnswer(query, top_documents, prompt_model, timeout_seconds: int = 30):
@@ -28,6 +59,8 @@ def GenerateAnswer(query, top_documents, prompt_model, timeout_seconds: int = 30
         Answer:
         """
 
+        Enforce_token_limit(prompt)
+
         # Call Groq API (Llama 3.3-70B)
         completion = groq_client.chat.completions.create(
             model=prompt_model,

requirements.txt

@@ -5,4 +5,5 @@ huggingface_hub
 pymilvus
 nltk
 sentence-transformers
-Groq
+Groq
+tiktoken