Add robust response validation and token config support

Introduces consistent None and truncation checks for model responses across all major modules, improving error handling and fallback behavior. Adds support for max_completion_tokens in request configuration, with priority over max_tokens, and propagates this through plugin and server layers. Enhances logging for truncated or empty responses and ensures backward compatibility for token limit parameters.

Author: codelion

optillm/bon.py

@@ -22,13 +22,17 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
             "temperature": 1
         }
         response = client.chat.completions.create(**provider_request)
-        
+
         # Log provider call
         if request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-        
-        completions = [choice.message.content for choice in response.choices]
+
+        # Check for valid response with None-checking
+        if response is None or not response.choices:
+            raise Exception("Response is None or has no choices")
+
+        completions = [choice.message.content for choice in response.choices if choice.message.content is not None]
         logger.info(f"Generated {len(completions)} initial completions using n parameter. Tokens used: {response.usage.completion_tokens}")
         bon_completion_tokens += response.usage.completion_tokens
 
@@ -46,12 +50,20 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
                     "temperature": 1
                 }
                 response = client.chat.completions.create(**provider_request)
-                
+
                 # Log provider call
                 if request_id:
                     response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
                     conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-                
+
+                # Check for valid response with None-checking
+                if (response is None or
+                    not response.choices or
+                    response.choices[0].message.content is None or
+                    response.choices[0].finish_reason == "length"):
+                    logger.warning(f"Completion {i+1}/{n} truncated or empty, skipping")
+                    continue
+
                 completions.append(response.choices[0].message.content)
                 bon_completion_tokens += response.usage.completion_tokens
                 logger.debug(f"Generated completion {i+1}/{n}")
@@ -83,18 +95,27 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
             "temperature": 0.1
         }
         rating_response = client.chat.completions.create(**provider_request)
-        
+
         # Log provider call
         if request_id:
             response_dict = rating_response.model_dump() if hasattr(rating_response, 'model_dump') else rating_response
             conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-        
+
         bon_completion_tokens += rating_response.usage.completion_tokens
-        try:
-            rating = float(rating_response.choices[0].message.content.strip())
-            ratings.append(rating)
-        except ValueError:
+
+        # Check for valid response with None-checking
+        if (rating_response is None or
+            not rating_response.choices or
+            rating_response.choices[0].message.content is None or
+            rating_response.choices[0].finish_reason == "length"):
+            logger.warning("Rating response truncated or empty, using default rating of 0")
             ratings.append(0)
+        else:
+            try:
+                rating = float(rating_response.choices[0].message.content.strip())
+                ratings.append(rating)
+            except ValueError:
+                ratings.append(0)
 
         rating_messages = rating_messages[:-2]
 

optillm/mcts.py

@@ -122,13 +122,18 @@ def generate_actions(self, state: DialogueState) -> List[str]:
             "temperature": 1
         }
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call
         if self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
-        
-        completions = [choice.message.content.strip() for choice in response.choices]
+
+        # Check for valid response with None-checking
+        if response is None or not response.choices:
+            logger.error("Failed to get valid completions from the model")
+            return []
+
+        completions = [choice.message.content.strip() for choice in response.choices if choice.message.content is not None]
         self.completion_tokens += response.usage.completion_tokens
         logger.info(f"Received {len(completions)} completions from the model")
         return completions
@@ -151,13 +156,22 @@ def apply_action(self, state: DialogueState, action: str) -> DialogueState:
             "temperature": 1
         }
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call
         if self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
-        
-        next_query = response.choices[0].message.content
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.warning("Next query response truncated or empty, using default")
+            next_query = "Please continue."
+        else:
+            next_query = response.choices[0].message.content
+
         self.completion_tokens += response.usage.completion_tokens
         logger.info(f"Generated next user query: {next_query}")
         return DialogueState(state.system_prompt, new_history, next_query)
@@ -181,13 +195,22 @@ def evaluate_state(self, state: DialogueState) -> float:
             "temperature": 0.1
         }
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call
         if self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
-        
+
         self.completion_tokens += response.usage.completion_tokens
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.warning("Evaluation response truncated or empty. Using default value 0.5")
+            return 0.5
+
         try:
             score = float(response.choices[0].message.content.strip())
             score = max(0, min(score, 1))  # Ensure the score is between 0 and 1

optillm/moa.py

@@ -25,15 +25,19 @@ def mixture_of_agents(system_prompt: str, initial_query: str, client, model: str
         }
 
         response = client.chat.completions.create(**provider_request)
-        
+
         # Convert response to dict for logging
         response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
-        
+
         # Log provider call if conversation logging is enabled
         if request_id:
             conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-        
-        completions = [choice.message.content for choice in response.choices]
+
+        # Check for valid response with None-checking
+        if response is None or not response.choices:
+            raise Exception("Response is None or has no choices")
+
+        completions = [choice.message.content for choice in response.choices if choice.message.content is not None]
         moa_completion_tokens += response.usage.completion_tokens
         logger.info(f"Generated {len(completions)} initial completions using n parameter. Tokens used: {response.usage.completion_tokens}")
 
@@ -56,14 +60,22 @@ def mixture_of_agents(system_prompt: str, initial_query: str, client, model: str
                 }
 
                 response = client.chat.completions.create(**provider_request)
-                
+
                 # Convert response to dict for logging
                 response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
-                
+
                 # Log provider call if conversation logging is enabled
                 if request_id:
                     conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-                
+
+                # Check for valid response with None-checking
+                if (response is None or
+                    not response.choices or
+                    response.choices[0].message.content is None or
+                    response.choices[0].finish_reason == "length"):
+                    logger.warning(f"Completion {i+1}/3 truncated or empty, skipping")
+                    continue
+
                 completions.append(response.choices[0].message.content)
                 moa_completion_tokens += response.usage.completion_tokens
                 logger.debug(f"Generated completion {i+1}/3")
@@ -118,15 +130,24 @@ def mixture_of_agents(system_prompt: str, initial_query: str, client, model: str
     }
 
     critique_response = client.chat.completions.create(**provider_request)
-    
+
     # Convert response to dict for logging
     response_dict = critique_response.model_dump() if hasattr(critique_response, 'model_dump') else critique_response
-    
+
     # Log provider call if conversation logging is enabled
     if request_id:
         conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-    
-    critiques = critique_response.choices[0].message.content
+
+    # Check for valid response with None-checking
+    if (critique_response is None or
+        not critique_response.choices or
+        critique_response.choices[0].message.content is None or
+        critique_response.choices[0].finish_reason == "length"):
+        logger.warning("Critique response truncated or empty, using generic critique")
+        critiques = "All candidates show reasonable approaches to the problem."
+    else:
+        critiques = critique_response.choices[0].message.content
+
     moa_completion_tokens += critique_response.usage.completion_tokens
     logger.info(f"Generated critiques. Tokens used: {critique_response.usage.completion_tokens}")
 
@@ -165,16 +186,27 @@ def mixture_of_agents(system_prompt: str, initial_query: str, client, model: str
     }
 
     final_response = client.chat.completions.create(**provider_request)
-    
+
     # Convert response to dict for logging
     response_dict = final_response.model_dump() if hasattr(final_response, 'model_dump') else final_response
-    
+
     # Log provider call if conversation logging is enabled
     if request_id:
         conversation_logger.log_provider_call(request_id, provider_request, response_dict)
-    
+
     moa_completion_tokens += final_response.usage.completion_tokens
     logger.info(f"Generated final response. Tokens used: {final_response.usage.completion_tokens}")
-    
+
+    # Check for valid response with None-checking
+    if (final_response is None or
+        not final_response.choices or
+        final_response.choices[0].message.content is None or
+        final_response.choices[0].finish_reason == "length"):
+        logger.error("Final response truncated or empty. Consider increasing max_tokens.")
+        # Return best completion if final response failed
+        result = completions[0] if completions else "Error: Response was truncated due to token limit. Please increase max_tokens or max_completion_tokens."
+    else:
+        result = final_response.choices[0].message.content
+
     logger.info(f"Total completion tokens used: {moa_completion_tokens}")
-    return final_response.choices[0].message.content, moa_completion_tokens
+    return result, moa_completion_tokens

optillm/plansearch.py

@@ -35,12 +35,21 @@ def generate_observations(self, problem: str, num_observations: int = 3) -> List
         }
 
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call if conversation logging is enabled
         if hasattr(optillm, 'conversation_logger') and optillm.conversation_logger and self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.plansearch_completion_tokens += response.usage.completion_tokens
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.warning("Observations response truncated or empty, returning empty list")
+            return []
+
         observations = response.choices[0].message.content.strip().split('\n')
         return [obs.strip() for obs in observations if obs.strip()]
 
@@ -70,12 +79,21 @@ def generate_derived_observations(self, problem: str, observations: List[str], n
         }
 
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call if conversation logging is enabled
         if hasattr(optillm, 'conversation_logger') and optillm.conversation_logger and self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.plansearch_completion_tokens += response.usage.completion_tokens
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.warning("Derived observations response truncated or empty, returning empty list")
+            return []
+
         new_observations = response.choices[0].message.content.strip().split('\n')
         return [obs.strip() for obs in new_observations if obs.strip()]
 
@@ -101,14 +119,23 @@ def generate_solution(self, problem: str, observations: List[str]) -> str:
                 {"role": "user", "content": prompt}
             ]
         }
-        
+
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call if conversation logging is enabled
         if hasattr(optillm, 'conversation_logger') and optillm.conversation_logger and self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.plansearch_completion_tokens += response.usage.completion_tokens
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.error("Solution generation response truncated or empty. Consider increasing max_tokens.")
+            return "Error: Response was truncated due to token limit. Please increase max_tokens or max_completion_tokens."
+
         return response.choices[0].message.content.strip()
 
     def implement_solution(self, problem: str, solution: str) -> str:
@@ -134,14 +161,23 @@ def implement_solution(self, problem: str, solution: str) -> str:
                 {"role": "user", "content": prompt}
             ]
         }
-        
+
         response = self.client.chat.completions.create(**provider_request)
-        
+
         # Log provider call if conversation logging is enabled
         if hasattr(optillm, 'conversation_logger') and optillm.conversation_logger and self.request_id:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.plansearch_completion_tokens += response.usage.completion_tokens
+
+        # Check for valid response with None-checking
+        if (response is None or
+            not response.choices or
+            response.choices[0].message.content is None or
+            response.choices[0].finish_reason == "length"):
+            logger.error("Implementation response truncated or empty. Consider increasing max_tokens.")
+            return "Error: Response was truncated due to token limit. Please increase max_tokens or max_completion_tokens."
+
         return response.choices[0].message.content.strip()
 
     def solve(self, problem: str, num_initial_observations: int = 3, num_derived_observations: int = 2) -> Tuple[str, str]: