Update agent_logic.py

agent_logic.py

@@ -1,4 +1,4 @@
-# agent_logic.py (Now reads prompts from config)
+# agent_logic.py (Now with error logging)
 import asyncio
 from typing import AsyncGenerator, Dict, Optional
 import json
@@ -7,18 +7,15 @@ import google.generativeai as genai
 from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
 from personas import PERSONAS_DATA 
-
-# 1. Import all settings from our new config files
 import config
 from utils import load_prompt
 from mcp_servers import AgentCalibrator, BusinessSolutionEvaluator, get_llm_response
 
-# 2. Load prompts from files AT STARTUP
 CLASSIFIER_SYSTEM_PROMPT = load_prompt(config.PROMPT_FILES["classifier"])
 HOMOGENEOUS_MANAGER_PROMPT = load_prompt(config.PROMPT_FILES["manager_homogeneous"])
 HETEROGENEOUS_MANAGER_PROMPT = load_prompt(config.PROMPT_FILES["manager_heterogeneous"])
 
-# --- 3. Define the Specialist Agent Architectures ---
+# --- (Specialist Agent classes are unchanged) ---
 class Baseline_Single_Agent:
     def __init__(self, api_clients: dict):
         self.gemini_client = api_clients.get("Gemini")
@@ -32,7 +29,7 @@ class Baseline_Single_Agent:
 class Baseline_Static_Homogeneous:
     def __init__(self, api_clients: dict):
         self.api_clients = {name: client for name, client in api_clients.items() if client}
-        self.gemini_client = api_clients.get("Gemini") # Manager is always Gemini
+        self.gemini_client = api_clients.get("Gemini")
 
     async def solve(self, problem: str, persona_prompt: str):
         print(f"--- (Specialist Team: Homogeneous) solving (live)... ---")
@@ -48,7 +45,6 @@ class Baseline_Static_Homogeneous:
             
         responses = await asyncio.gather(*tasks)
         
-        # 4. Use the loaded manager prompt
         manager_system_prompt = HOMOGENEOUS_MANAGER_PROMPT
         reports_str = "\n\n".join(f"Report from Team Member {i+1}:\n{resp}" for i, resp in enumerate(responses))
         manager_user_prompt = f"Original Problem: {problem}\n\n{reports_str}\n\nPlease synthesize these reports into one final, comprehensive solution."
@@ -58,7 +54,7 @@ class Baseline_Static_Homogeneous:
 class Baseline_Static_Heterogeneous:
     def __init__(self, api_clients: dict):
         self.api_clients = api_clients
-        self.gemini_client = api_clients.get("Gemini") # Manager is always Gemini
+        self.gemini_client = api_clients.get("Gemini")
         
     async def solve(self, problem: str, team_plan: dict):
         print(f"--- (Specialist Team: Heterogeneous) solving (live)... ---")
@@ -68,7 +64,7 @@ class Baseline_Static_Heterogeneous:
         tasks = []
         for role, config_data in team_plan.items():
             llm_name = config_data["llm"]
-            persona_key = config_data["persona"] # e.g., "Culture_5"
+            persona_key = config_data["persona"]
             client = self.api_clients.get(llm_name)
             
             if not client:
@@ -82,7 +78,6 @@ class Baseline_Static_Heterogeneous:
             
         responses = await asyncio.gather(*tasks)
 
-        # 5. Use the loaded manager prompt
         manager_system_prompt = HETEROGENEOUS_MANAGER_PROMPT
         reports_str = "\n\n".join(f"Report from {team_plan[role]['llm']} (as {role}):\n{resp}" for (role, resp) in zip(team_plan.keys(), responses))
         manager_user_prompt = f"Original Problem: {problem}\n\n{reports_str}\n\nPlease synthesize these specialist reports into one final, comprehensive solution."
@@ -94,17 +89,11 @@ class StrategicSelectorAgent:
     
     def __init__(self, api_keys: Dict[str, Optional[str]]):
         self.api_keys = api_keys
-        
-        self.api_clients = {
-            "Gemini": None,
-            "Anthropic": None,
-            "SambaNova": None
-        }
+        self.api_clients = { "Gemini": None, "Anthropic": None, "SambaNova": None }
         
         if api_keys.get("google"):
             try:
                 genai.configure(api_key=api_keys["google"])
-                # 6. Read model name from config
                 self.api_clients["Gemini"] = genai.GenerativeModel(config.MODELS["Gemini"]["default"])
             except Exception as e:
                 print(f"Warning: Failed to initialize Gemini client. Error: {e}")
@@ -118,10 +107,7 @@ class StrategicSelectorAgent:
         if api_keys.get("sambanova"):
             try:
                 base_url = os.getenv("SAMBANOVA_BASE_URL", "https://api.sambanova.ai/v1")
-                self.api_clients["SambaNova"] = AsyncOpenAI(
-                    api_key=api_keys["sambanova"],
-                    base_url=base_url
-                )
+                self.api_clients["SambaNova"] = AsyncOpenAI(api_key=api_keys["sambanova"], base_url=base_url)
             except Exception as e:
                 print(f"Warning: Failed to initialize SambaNova client. Error: {e}")
         
@@ -135,14 +121,11 @@ class StrategicSelectorAgent:
         self.homo_team = Baseline_Static_Homogeneous(self.api_clients)
         self.hetero_team = Baseline_Static_Heterogeneous(self.api_clients)
 
-        # 7. Check if prompts loaded correctly
         if "ERROR:" in CLASSIFIER_SYSTEM_PROMPT:
             raise FileNotFoundError(CLASSIFIER_SYSTEM_PROMPT)
 
     async def _classify_problem(self, problem: str) -> AsyncGenerator[str, None]:
         yield "Classifying problem archetype (live)..."
-        
-        # 8. Use the loaded classifier prompt
         classification = await get_llm_response(
             "Gemini", 
             self.api_clients["Gemini"], 
@@ -167,7 +150,6 @@ class StrategicSelectorAgent:
         solution_draft = "" 
         
         try:
-            # 9. Read default persona from config
             default_persona = PERSONAS_DATA[config.DEFAULT_PERSONA_KEY]["description"]
 
             if classification == "Direct_Procedure" or classification == "Holistic_Abstract_Reasoning":
@@ -180,7 +162,17 @@ class StrategicSelectorAgent:
                 
             elif classification == "Cognitive_Labyrinth":
                 yield "Deploying: Static Heterogeneous Team (Cognitive Diversity)..."
-                team_plan = await self.calibrator.calibrate_team(problem)
+                
+                # --- NEW: Capture errors from calibration ---
+                team_plan, calibration_errors = await self.calibrator.calibrate_team(problem)
+                
+                # --- NEW: Yield any calibration errors ---
+                if calibration_errors:
+                    yield "--- CALIBRATION WARNINGS ---"
+                    for err in calibration_errors:
+                        yield err
+                    yield "-----------------------------"
+                    
                 yield f"Calibration complete. Best Team: {json.dumps({k: v['llm'] for k, v in team_plan.items()})}"
                 solution_draft = await self.hetero_team.solve(problem, team_plan)
             
@@ -199,10 +191,11 @@ class StrategicSelectorAgent:
             scores = {k: v.get('score', 0) for k, v in v_fitness_json.items()}
             yield f"Initial Score: {scores}"
             
-            yield "Skipping self-correction for Milestone 4..."
+            # --- This is where Milestone 5 will go ---
+            yield "Skipping self-correction for now..."
             
             await asyncio.sleep(0.5)
-            yield "Milestone 4 Complete. All systems are live and stateless."
+            yield "Milestone 4 (with error logging) Complete."
             
             solution_draft_json_safe = json.dumps(solution_draft)
             yield f"FINAL: {{\"text\": {solution_draft_json_safe}, \"audio\": null}}"