Implement conservative token limits and better context management

aimakerspace/openai_utils/chatmodel.py

@@ -11,8 +11,10 @@ class ChatOpenAI:
         self.openai_api_key = os.getenv("OPENAI_API_KEY")
         if self.openai_api_key is None:
             raise ValueError("OPENAI_API_KEY is not set")
-        self.max_tokens = 8192  # Maximum tokens for response
-        self.max_total_tokens = 16384  # Maximum total tokens (prompt + response)
+        # More conservative token limits
+        self.max_tokens = 4000  # Reduced from 8192 to leave room for context
+        self.max_total_tokens = 8000  # Reduced from 16384 to be safe
+        self.temperature = 0.7
 
     def run(self, messages, text_only: bool = True, **kwargs):
         if not isinstance(messages, list):
@@ -20,11 +22,23 @@ class ChatOpenAI:
 
         client = OpenAI()
         try:
+            # Estimate tokens in messages (rough estimate: 1 token ≈ 4 characters)
+            total_chars = sum(len(str(msg.get('content', ''))) for msg in messages)
+            estimated_tokens = total_chars // 4
+            
+            if estimated_tokens > 4000:  # If messages are too long
+                print(f"Warning: Messages too long ({estimated_tokens} estimated tokens). Truncating...")
+                # Keep only the most recent messages that fit
+                while estimated_tokens > 4000 and len(messages) > 2:  # Keep system message and last user message
+                    messages.pop(1)  # Remove oldest message after system message
+                    total_chars = sum(len(str(msg.get('content', ''))) for msg in messages)
+                    estimated_tokens = total_chars // 4
+
             response = client.chat.completions.create(
                 model=self.model_name,
                 messages=messages,
                 max_tokens=self.max_tokens,
-                temperature=0.7,  # Add some creativity while maintaining accuracy
+                temperature=self.temperature,
                 **kwargs
             )
             if text_only:
@@ -41,11 +55,23 @@ class ChatOpenAI:
         client = AsyncOpenAI()
 
         try:
+            # Estimate tokens in messages (rough estimate: 1 token ≈ 4 characters)
+            total_chars = sum(len(str(msg.get('content', ''))) for msg in messages)
+            estimated_tokens = total_chars // 4
+            
+            if estimated_tokens > 4000:  # If messages are too long
+                print(f"Warning: Messages too long ({estimated_tokens} estimated tokens). Truncating...")
+                # Keep only the most recent messages that fit
+                while estimated_tokens > 4000 and len(messages) > 2:  # Keep system message and last user message
+                    messages.pop(1)  # Remove oldest message after system message
+                    total_chars = sum(len(str(msg.get('content', ''))) for msg in messages)
+                    estimated_tokens = total_chars // 4
+
             stream = await client.chat.completions.create(
                 model=self.model_name,
                 messages=messages,
                 max_tokens=self.max_tokens,
-                temperature=0.7,  # Add some creativity while maintaining accuracy
+                temperature=self.temperature,
                 stream=True,
                 **kwargs
             )

app.py

@@ -32,12 +32,12 @@ class RetrievalAugmentedQAPipeline:
 
     async def arun_pipeline(self, user_query: str):
         # Get more contexts but limit the total length
-        context_list = self.vector_db_retriever.search_by_text(user_query, k=6)
+        context_list = self.vector_db_retriever.search_by_text(user_query, k=3)  # Reduced from 6 to 3
         
-        # Limit total context length to approximately 6000 tokens (24000 characters)
+        # Limit total context length to approximately 3000 tokens (12000 characters)
         context_prompt = ""
         total_length = 0
-        max_length = 24000  # Rough estimate: 1 token ≈ 4 characters
+        max_length = 12000  # Reduced from 24000 to 12000
         
         for context in context_list:
             if total_length + len(context[0]) > max_length: