Add fallback for providers without n parameter support

Updated best_of_n_sampling, mixture_of_agents, and majority_voting_plugin to handle providers that do not support the 'n' parameter by generating completions/candidates one by one in a loop. This improves compatibility with a wider range of API providers and ensures robust completion generation even when batch generation is not available.

Author: codelion

optillm/bon.py

@@ -10,16 +10,45 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
 
     completions = []
 
-    response = client.chat.completions.create(
-        model=model,
-        messages=messages,
-        max_tokens=4096,
-        n=n,
-        temperature=1
-    )
-    completions = [choice.message.content for choice in response.choices]
-    logger.info(f"Generated {len(completions)} initial completions. Tokens used: {response.usage.completion_tokens}")
-    bon_completion_tokens += response.usage.completion_tokens
+    try:
+        # Try to generate n completions in a single API call using n parameter
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            max_tokens=4096,
+            n=n,
+            temperature=1
+        )
+        completions = [choice.message.content for choice in response.choices]
+        logger.info(f"Generated {len(completions)} initial completions using n parameter. Tokens used: {response.usage.completion_tokens}")
+        bon_completion_tokens += response.usage.completion_tokens
+        
+    except Exception as e:
+        logger.warning(f"n parameter not supported by provider: {str(e)}")
+        logger.info(f"Falling back to generating {n} completions one by one")
+        
+        # Fallback: Generate completions one by one in a loop
+        for i in range(n):
+            try:
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=messages,
+                    max_tokens=4096,
+                    temperature=1
+                )
+                completions.append(response.choices[0].message.content)
+                bon_completion_tokens += response.usage.completion_tokens
+                logger.debug(f"Generated completion {i+1}/{n}")
+                
+            except Exception as fallback_error:
+                logger.error(f"Error generating completion {i+1}: {str(fallback_error)}")
+                continue
+        
+        if not completions:
+            logger.error("Failed to generate any completions")
+            return "Error: Could not generate any completions", 0
+        
+        logger.info(f"Generated {len(completions)} completions using fallback method. Total tokens used: {bon_completion_tokens}")
 
     # Rate the completions
     rating_messages = messages.copy()

optillm/moa.py

@@ -8,19 +8,61 @@ def mixture_of_agents(system_prompt: str, initial_query: str, client, model: str
     completions = []
 
     logger.debug(f"Generating initial completions for query: {initial_query}")
-    response = client.chat.completions.create(
-        model=model,
-        messages=[
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": initial_query}
-        ],
-        max_tokens=4096,
-        n=3,
-        temperature=1
-    )
-    completions = [choice.message.content for choice in response.choices]
-    moa_completion_tokens += response.usage.completion_tokens
-    logger.info(f"Generated {len(completions)} initial completions. Tokens used: {response.usage.completion_tokens}")
+    
+    try:
+        # Try to generate 3 completions in a single API call using n parameter
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": initial_query}
+            ],
+            max_tokens=4096,
+            n=3,
+            temperature=1
+        )
+        completions = [choice.message.content for choice in response.choices]
+        moa_completion_tokens += response.usage.completion_tokens
+        logger.info(f"Generated {len(completions)} initial completions using n parameter. Tokens used: {response.usage.completion_tokens}")
+        
+    except Exception as e:
+        logger.warning(f"n parameter not supported by provider: {str(e)}")
+        logger.info("Falling back to generating 3 completions one by one")
+        
+        # Fallback: Generate 3 completions one by one in a loop
+        completions = []
+        for i in range(3):
+            try:
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": initial_query}
+                    ],
+                    max_tokens=4096,
+                    temperature=1
+                )
+                completions.append(response.choices[0].message.content)
+                moa_completion_tokens += response.usage.completion_tokens
+                logger.debug(f"Generated completion {i+1}/3")
+                
+            except Exception as fallback_error:
+                logger.error(f"Error generating completion {i+1}: {str(fallback_error)}")
+                continue
+        
+        if not completions:
+            logger.error("Failed to generate any completions")
+            return "Error: Could not generate any completions", 0
+        
+        logger.info(f"Generated {len(completions)} completions using fallback method. Total tokens used: {moa_completion_tokens}")
+    
+    # Handle case where fewer than 3 completions were generated
+    if len(completions) < 3:
+        original_count = len(completions)
+        # Pad with the first completion to ensure we have 3
+        while len(completions) < 3:
+            completions.append(completions[0])
+        logger.warning(f"Only generated {original_count} unique completions, padded to 3 for critique")
 
     logger.debug("Preparing critique prompt")
     critique_prompt = f"""

optillm/plugins/majority_voting_plugin.py

@@ -213,65 +213,81 @@ def run(
         candidates = [choice.message.content for choice in response.choices]
         total_tokens = response.usage.completion_tokens
 
-        logger.info(f"Generated {len(candidates)} candidates. Tokens used: {total_tokens}")
+        logger.info(f"Generated {len(candidates)} candidates using n parameter. Tokens used: {total_tokens}")
 
-        # Extract answers from each candidate
-        answers = []
-        answer_to_response = {}  # Map normalized answers to full responses
-        
-        for i, candidate in enumerate(candidates):
-            answer = extract_answer(candidate)
-            if answer:
-                normalized = normalize_answer(answer)
-                answers.append(normalized)
-                # Keep the first full response for each unique answer
-                if normalized not in answer_to_response:
-                    answer_to_response[normalized] = candidate
-                logger.debug(f"Candidate {i+1} answer: {answer} (normalized: {normalized})")
-            else:
-                logger.warning(f"Could not extract answer from candidate {i+1}")
-        
-        if not answers:
-            logger.warning("No answers could be extracted from any candidate")
-            # Return the first candidate as fallback
-            return candidates[0] if candidates else "Error: No candidates generated", total_tokens
-        
-        # Count answer frequencies
-        answer_counts = Counter(answers)
-        logger.info(f"Answer distribution: {dict(answer_counts)}")
-        
-        # Get the most common answer
-        most_common_answer, count = answer_counts.most_common(1)[0]
-        confidence = count / len(answers)
-        
-        logger.info(f"Most common answer: '{most_common_answer}' with {count}/{len(answers)} votes ({confidence:.1%} confidence)")
-        
-        # Get the full response corresponding to the most common answer
-        winning_response = answer_to_response.get(most_common_answer, candidates[0])
-        
-        # Log voting summary to console instead of adding to response
-        logger.info("Majority Voting Summary:")
-        logger.info(f"  - Generated {k} candidates")
-        logger.info(f"  - Most common answer: {most_common_answer}")
-        logger.info(f"  - Votes: {count}/{len(answers)} ({confidence:.1%} confidence)")
-        
-        if len(answer_counts) > 1:
-            other_answers = [f"{ans} ({cnt} votes)" for ans, cnt in answer_counts.items() if ans != most_common_answer]
-            logger.info(f"  - Other answers: {', '.join(other_answers)}")
+    except Exception as e:
+        logger.warning(f"n parameter not supported by provider: {str(e)}")
+        logger.info(f"Falling back to generating {k} candidates one by one")
 
-        # Return only the full response from the winning answer
-        return winning_response, total_tokens
+        # Fallback: Generate candidates one by one in a loop
+        candidates = []
+        total_tokens = 0
 
-    except Exception as e:
-        logger.error(f"Error in majority voting: {str(e)}")
-        # Fall back to single response
-        logger.info("Falling back to single response generation")
+        for i in range(k):
+            try:
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+                candidates.append(response.choices[0].message.content)
+                total_tokens += response.usage.completion_tokens
+                logger.debug(f"Generated candidate {i+1}/{k}")
+                
+            except Exception as fallback_error:
+                logger.error(f"Error generating candidate {i+1}: {str(fallback_error)}")
+                continue
 
-        response = client.chat.completions.create(
-            model=model,
-            messages=messages,
-            temperature=temperature,
-            max_tokens=max_tokens
-        )
+        if not candidates:
+            logger.error("Failed to generate any candidates")
+            return "Error: Could not generate any candidates", 0
 
-        return response.choices[0].message.content, response.usage.completion_tokens
+        logger.info(f"Generated {len(candidates)} candidates using fallback method. Total tokens used: {total_tokens}")
+    
+    # Extract answers from each candidate
+    answers = []
+    answer_to_response = {}  # Map normalized answers to full responses
+    
+    for i, candidate in enumerate(candidates):
+        answer = extract_answer(candidate)
+        if answer:
+            normalized = normalize_answer(answer)
+            answers.append(normalized)
+            # Keep the first full response for each unique answer
+            if normalized not in answer_to_response:
+                answer_to_response[normalized] = candidate
+            logger.debug(f"Candidate {i+1} answer: {answer} (normalized: {normalized})")
+        else:
+            logger.warning(f"Could not extract answer from candidate {i+1}")
+    
+    if not answers:
+        logger.warning("No answers could be extracted from any candidate")
+        # Return the first candidate as fallback
+        return candidates[0] if candidates else "Error: No candidates generated", total_tokens
+    
+    # Count answer frequencies
+    answer_counts = Counter(answers)
+    logger.info(f"Answer distribution: {dict(answer_counts)}")
+    
+    # Get the most common answer
+    most_common_answer, count = answer_counts.most_common(1)[0]
+    confidence = count / len(answers)
+    
+    logger.info(f"Most common answer: '{most_common_answer}' with {count}/{len(answers)} votes ({confidence:.1%} confidence)")
+    
+    # Get the full response corresponding to the most common answer
+    winning_response = answer_to_response.get(most_common_answer, candidates[0])
+    
+    # Log voting summary to console instead of adding to response
+    logger.info("Majority Voting Summary:")
+    logger.info(f"  - Generated {len(candidates)} candidates")
+    logger.info(f"  - Most common answer: {most_common_answer}")
+    logger.info(f"  - Votes: {count}/{len(answers)} ({confidence:.1%} confidence)")
+    
+    if len(answer_counts) > 1:
+        other_answers = [f"{ans} ({cnt} votes)" for ans, cnt in answer_counts.items() if ans != most_common_answer]
+        logger.info(f"  - Other answers: {', '.join(other_answers)}")
+    
+    # Return only the full response from the winning answer
+    return winning_response, total_tokens