Update mcp_servers.py

mcp_servers.py

@@ -1,30 +1,11 @@
-# mcp_servers.py (Corrected for GOOGLE_API_KEY)
+# mcp_servers.py (New Secure & Robust Version)
 import asyncio
-import os
-import httpx
 import json
-import google.generativeai as genai
-import anthropic
-import openai
+from typing import Dict, Optional
 from personas import PERSONAS_DATA
-
-# --- 1. Load API Keys from Blaxel Secrets ---
-# --- THIS IS THE FIX ---
-GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY") # Use the secret name from your screenshot
-# ---
-ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
-SAMBANOVA_API_KEY = os.getenv("SAMBANOVA_API_KEY")
-SAMBANOVA_BASE_URL = os.getenv("SAMBANOVA_BASE_URL", "https://api.sambanova.ai/v1")
-
-# --- 2. Configure API Clients ---
-genai.configure(api_key=GEMINI_API_KEY)
-gemini_model = genai.GenerativeModel('gemini-1.5-pro-latest')
-anthropic_client = anthropic.AsyncAnthropic(api_key=ANTHROPIC_API_KEY)
-
-sambanova_client = openai.AsyncOpenAI(
-    api_key=SAMBANOVA_API_KEY,
-    base_url=SAMBANOVA_BASE_URL
-)
+import google.generativeai as genai
+from anthropic import AsyncAnthropic
+from openai import AsyncOpenAI
 
 # This is the prompt from your 'LLM judges prompt v3.0.docx'
 EVALUATION_PROMPT_TEMPLATE = """
@@ -63,26 +44,30 @@ You MUST return *only* a valid JSON object in the following format:
 class BusinessSolutionEvaluator:
     """Implements the "LLM-as-a-Judge" with a live call to Gemini."""
     
+    def __init__(self, gemini_client: Optional[genai.GenerativeModel]):
+        # 1. Get the pre-initialized client from the main agent
+        if not gemini_client:
+            raise ValueError("BusinessSolutionEvaluator requires a Google/Gemini client to function.")
+        self.gemini_model = gemini_client
+    
     async def evaluate(self, problem: str, solution_text: str) -> dict:
         print(f"Evaluating solution (live): {solution_text[:50]}...")
-        
         prompt = EVALUATION_PROMPT_TEMPLATE.format(problem=problem, solution_text=solution_text)
         
         try:
-            response = await gemini_model.generate_content_async(
+            response = await self.gemini_model.generate_content_async(
                 prompt,
                 generation_config=genai.types.GenerationConfig(
                     response_mime_type="application/json"
                 )
             )
-            
             json_text = response.text.strip().replace("```json", "").replace("```", "")
             v_fitness = json.loads(json_text)
-            
             print(f"Evaluation complete (live): {v_fitness}")
             return v_fitness
         except Exception as e:
             print(f"ERROR: BusinessSolutionEvaluator failed: {e}")
+            # Fallback in case of API error
             return {
                 "Novelty": {"score": 1, "justification": "Error during evaluation."},
                 "Usefulness_Feasibility": {"score": 1, "justification": "Error during evaluation."},
@@ -91,52 +76,29 @@ class BusinessSolutionEvaluator:
                 "Cultural_Appropriateness": {"score": 1, "justification": "Error during evaluation."}
             }
 
-# --- 3. Unified API Call Function ---
-async def get_llm_response(client_name: str, system_prompt: str, user_prompt: str) -> str:
-    """A single function to handle calling any of the three sponsor LLMs."""
-    try:
-        if client_name == "Gemini":
-            chat = gemini_model.start_chat(history=[
-                {'role': 'user', 'parts': [system_prompt]},
-                {'role': 'model', 'parts': ["Understood. I will act as this persona."]}
-            ])
-            response = await chat.send_message_async(user_prompt)
-            return response.text
-
-        elif client_name == "Anthropic":
-            response = await anthropic_client.messages.create(
-                model="claude-3-opus-20240229",
-                max_tokens=2048,
-                system=system_prompt,
-                messages=[{"role": "user", "content": user_prompt}]
-            )
-            return response.content[0].text
-        
-        elif client_name == "SambaNova":
-            completion = await sambanova_client.chat.completions.create(
-                model="Meta-Llama-3.1-8B-Instruct",
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt}
-                ]
-            )
-            return completion.choices[0].message.content
-            
-    except Exception as e:
-        print(f"ERROR: API call to {client_name} failed: {e}")
-        return f"Error generating response from {client_name}."
-
-
 class AgentCalibrator:
-    """Tests the sponsor LLMs with live API calls."""
+    """Tests the available sponsor LLMs to find the best one for each role."""
     
-    def __init__(self, evaluator: BusinessSolutionEvaluator):
+    def __init__(self, api_clients: dict, evaluator: BusinessSolutionEvaluator):
         self.evaluator = evaluator
-        self.sponsor_llms = ["Gemini", "Anthropic", "SambaNova"]
+        # 1. Get the dict of *live* clients
+        self.api_clients = {name: client for name, client in api_clients.items() if client}
+        self.sponsor_llms = list(self.api_clients.keys())
+        print(f"AgentCalibrator initialized with enabled clients: {self.sponsor_llms}")
         
     async def calibrate_team(self, problem: str) -> dict:
-        print("Running LIVE calibration test for specialist team...")
+        print(f"Running LIVE calibration test for specialist team on {self.sponsor_llms}...")
         
+        # If only one LLM is available, skip calibration and default to it
+        if len(self.sponsor_llms) == 1:
+            default_llm = self.sponsor_llms[0]
+            print("Only one LLM available. Skipping calibration.")
+            return {
+                "Plant": {"persona": "Culture_5", "llm": default_llm},
+                "Implementer": {"persona": "Culture_Expert", "llm": default_llm},
+                "Monitor": {"persona": "Culture_11", "llm": default_llm}
+            }
+
         roles_to_test = {
             "Plant": PERSONAS_DATA["Culture_5"]["description"],
             "Implementer": PERSONAS_DATA["Culture_Expert"]["description"],
@@ -147,8 +109,8 @@ class AgentCalibrator:
         
         tasks = []
         for role, persona in roles_to_test.items():
-            for llm in self.sponsor_llms:
-                tasks.append(self.run_calibration_test(problem, role, llm, persona, test_problem))
+            for llm_name in self.sponsor_llms:
+                tasks.append(self.run_calibration_test(problem, role, llm_name, persona, test_problem))
         
         results = await asyncio.gather(*tasks)
         
@@ -161,7 +123,7 @@ class AgentCalibrator:
         
         for role in roles_to_test.keys():
             best_score = -1
-            best_llm = "None"
+            best_llm = self.sponsor_llms[0] # Default to first available
             for res in results:
                 if res["role"] == role:
                     metric = role_metrics[role]
@@ -180,16 +142,51 @@ class AgentCalibrator:
         print(f"Calibration complete (live). Team plan: {team_plan}")
         return team_plan
 
-    async def run_calibration_test(self, problem, role, llm, persona, test_problem):
+    async def run_calibration_test(self, problem, role, llm_name, persona, test_problem):
         """Helper to run a single test and evaluation."""
-        print(f"...Calibrating {role} on {llm}...")
-        solution = await get_llm_response(llm, persona, test_problem)
+        print(f"...Calibrating {role} on {llm_name}...")
+        client = self.api_clients[llm_name]
+        solution = await get_llm_response(llm_name, client, persona, test_problem)
+        
         if "Error generating response" in solution:
-            return {"role": role, "llM": llm, "score": {
-                "Novelty": {"score": 0}, 
-                "Usefulness_Feasibility": {"score": 0}, 
-                "Cultural_Appropriateness": {"score": 0}
+            return {"role": role, "llm": llm_name, "score": {
+                "Novelty": {"score": 0}, "Usefulness_Feasibility": {"score": 0}, "Cultural_Appropriateness": {"score": 0}
             }}
         
         score = await self.evaluator.evaluate(problem, solution)
-        return {"role": role, "llm": llm, "score": score}
+        return {"role": role, "llm": llm_name, "score": score}
+
+# --- Unified API Call Function ---
+async def get_llm_response(client_name: str, client, system_prompt: str, user_prompt: str) -> str:
+    """A single function to handle calling any of the three sponsor LLMs."""
+    try:
+        if client_name == "Gemini":
+            chat = client.start_chat(history=[
+                {'role': 'user', 'parts': [system_prompt]},
+                {'role': 'model', 'parts': ["Understood. I will act as this persona."]}
+            ])
+            response = await chat.send_message_async(user_prompt)
+            return response.text
+
+        elif client_name == "Anthropic":
+            response = await client.messages.create(
+                model="claude-3-opus-20240229",
+                max_tokens=2048,
+                system=system_prompt,
+                messages=[{"role": "user", "content": user_prompt}]
+            )
+            return response.content[0].text
+        
+        elif client_name == "SambaNova":
+            completion = await client.chat.completions.create(
+                model="Meta-Llama-3.1-8B-Instruct",
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ]
+            )
+            return completion.choices[0].message.content
+            
+    except Exception as e:
+        print(f"ERROR: API call to {client_name} failed: {e}")
+        return f"Error generating response from {client_name}."