Update mcp_servers.py

mcp_servers.py

@@ -1,4 +1,4 @@
-# mcp_servers.py (FIXED: Guarantees 4-value return tuple)
+# mcp_servers.py (FIXED: Added OpenAI & Nebius Support to get_llm_response)
 import asyncio
 import json
 import re
@@ -94,7 +94,6 @@ class AgentCalibrator:
         if not self.sponsor_llms:
              raise Exception("AgentCalibrator cannot run: No LLM clients are configured.")
         
-        # If only one model, return default plan + empty lists for details/usage
         if len(self.sponsor_llms) == 1:
             default_llm = self.sponsor_llms[0]
             print("Only one LLM available. Skipping calibration.")
@@ -103,7 +102,6 @@ class AgentCalibrator:
                 "Implementer": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Implementer"], "llm": default_llm},
                 "Monitor": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Monitor"], "llm": default_llm}
             }
-            # MUST RETURN 4 VALUES
             return plan, error_log, [], []
 
         roles_to_test = {role: PERSONAS_DATA[key]["description"] for role, key in config.CALIBRATION_CONFIG["roles_to_test"].items()}
@@ -154,7 +152,6 @@ class AgentCalibrator:
             "Monitor": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Monitor"], "llm": best_llms["Monitor"]}
         }
         print(f"Calibration complete (live). Team plan: {team_plan}")
-        # MUST RETURN 4 VALUES
         return team_plan, error_log, detailed_results, all_usage_stats
 
     async def run_calibration_test(self, problem, role, llm_name, persona, test_problem):
@@ -169,6 +166,7 @@ class AgentCalibrator:
             "role": role, "llm": llm_name, "score": score, "output": solution, "usage_gen": gen_usage, "usage_eval": eval_usage
         }
 
+# --- UPDATED: Handles OpenAI and Nebius ---
 async def get_llm_response(client_name: str, client, system_prompt: str, user_prompt: str) -> Tuple[str, dict]:
     """Returns (text_response, usage_dict)"""
     usage = {"model": client_name, "input": 0, "output": 0}
@@ -191,10 +189,18 @@ async def get_llm_response(client_name: str, client, system_prompt: str, user_pr
                 usage["output"] = response.usage.output_tokens
             return response.content[0].text, usage
         
+        # --- THIS IS THE PART THAT WAS MISSING OR INCOMPLETE ---
         elif client_name in ["SambaNova", "OpenAI", "Nebius"]:
+            
+            # Dynamically get the correct model ID from config.py
             model_id = config.MODELS.get(client_name, {}).get("default", "gpt-4o-mini")
+            
             completion = await client.chat.completions.create(
-                model=model_id, messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]
+                model=model_id,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ]
             )
             if hasattr(completion, "usage"):
                 usage["input"] = completion.usage.prompt_tokens