Update agent_logic.py

agent_logic.py

@@ -1,4 +1,4 @@
-# agent_logic.py (Now with error logging)
+# agent_logic.py (Milestone 5 - FINAL & ROBUST)
 import asyncio
 from typing import AsyncGenerator, Dict, Optional
 import json
@@ -10,83 +10,60 @@ from personas import PERSONAS_DATA
 import config
 from utils import load_prompt
 from mcp_servers import AgentCalibrator, BusinessSolutionEvaluator, get_llm_response
+from self_correction import SelfCorrector
+from async_generator import async_generator, yield_
 
 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"])
 
-# --- (Specialist Agent classes are unchanged) ---
+# (Baseline Agent Classes - UNCHANGED)
 class Baseline_Single_Agent:
     def __init__(self, api_clients: dict):
         self.gemini_client = api_clients.get("Gemini")
-
     async def solve(self, problem: str, persona_prompt: str):
-        print(f"--- (Specialist Team: Single Agent) solving (live)... ---")
-        if not self.gemini_client:
-            raise ValueError("Single_Agent requires a Google/Gemini client.")
+        if not self.gemini_client: raise ValueError("Single_Agent requires a Google/Gemini client.")
         return await get_llm_response("Gemini", self.gemini_client, persona_prompt, problem)
 
 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")
-
     async def solve(self, problem: str, persona_prompt: str):
-        print(f"--- (Specialist Team: Homogeneous) solving (live)... ---")
-        if not self.gemini_client:
-            raise ValueError("Homogeneous_Team requires a Google/Gemini client for its manager.")
-            
+        if not self.gemini_client: raise ValueError("Homogeneous_Team requires a Google/Gemini client.")
         system_prompt = persona_prompt
         user_prompt = f"As an expert Implementer, generate a detailed plan for this problem: {problem}"
-        
-        tasks = []
-        for llm_name, client in self.api_clients.items():
-            tasks.append(get_llm_response(llm_name, client, system_prompt, user_prompt))
-            
+        tasks = [get_llm_response(llm, client, system_prompt, user_prompt) for llm, client in self.api_clients.items()]
         responses = await asyncio.gather(*tasks)
-        
         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."
-        
         return await get_llm_response("Gemini", self.gemini_client, manager_system_prompt, manager_user_prompt)
 
 class Baseline_Static_Heterogeneous:
     def __init__(self, api_clients: dict):
         self.api_clients = api_clients
         self.gemini_client = api_clients.get("Gemini")
-        
     async def solve(self, problem: str, team_plan: dict):
-        print(f"--- (Specialist Team: Heterogeneous) solving (live)... ---")
-        if not self.gemini_client:
-            raise ValueError("Heterogeneous_Team requires a Google/Gemini client for its manager.")
-
+        if not self.gemini_client: raise ValueError("Heterogeneous_Team requires a Google/Gemini client.")
         tasks = []
         for role, config_data in team_plan.items():
             llm_name = config_data["llm"]
             persona_key = config_data["persona"]
             client = self.api_clients.get(llm_name)
-            
             if not client:
-                print(f"Warning: Calibrated LLM '{llm_name}' for role '{role}' is not available. Defaulting to Gemini.")
                 llm_name = "Gemini"
                 client = self.gemini_client
-            
             system_prompt = PERSONAS_DATA[persona_key]["description"]
             user_prompt = f"As the team's '{role}', provide your unique perspective on how to solve this problem: {problem}"
             tasks.append(get_llm_response(llm_name, client, system_prompt, user_prompt))
-            
         responses = await asyncio.gather(*tasks)
-
         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."
-        
         return await get_llm_response("Gemini", self.gemini_client, manager_system_prompt, manager_user_prompt)
 
 class StrategicSelectorAgent:
-    """This is MudabbirAI. It gets keys passed to it on creation."""
-    
     def __init__(self, api_keys: Dict[str, Optional[str]]):
         self.api_keys = api_keys
         self.api_clients = { "Gemini": None, "Anthropic": None, "SambaNova": None }
@@ -95,46 +72,114 @@ class StrategicSelectorAgent:
             try:
                 genai.configure(api_key=api_keys["google"])
                 self.api_clients["Gemini"] = genai.GenerativeModel(config.MODELS["Gemini"]["default"])
-            except Exception as e:
-                print(f"Warning: Failed to initialize Gemini client. Error: {e}")
-        
+            except Exception as e: print(f"Warning: Gemini init failed: {e}")
         if api_keys.get("anthropic"):
             try:
                 self.api_clients["Anthropic"] = AsyncAnthropic(api_key=api_keys["anthropic"])
-            except Exception as e:
-                print(f"Warning: Failed to initialize Anthropic client. Error: {e}")
-        
+            except Exception as e: print(f"Warning: Anthropic init failed: {e}")
         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)
-            except Exception as e:
-                print(f"Warning: Failed to initialize SambaNova client. Error: {e}")
+            except Exception as e: print(f"Warning: SambaNova init failed: {e}")
         
         if not self.api_clients["Gemini"]:
-            raise ValueError("Google API Key is required or invalid. The agent cannot function without its 'Judge'.")
+            raise ValueError("Google API Key is required.")
 
         self.evaluator = BusinessSolutionEvaluator(self.api_clients["Gemini"])
         self.calibrator = AgentCalibrator(self.api_clients, self.evaluator)
-        
+        self.corrector = SelfCorrector(threshold=3.0)
         self.single_agent = Baseline_Single_Agent(self.api_clients)
         self.homo_team = Baseline_Static_Homogeneous(self.api_clients)
         self.hetero_team = Baseline_Static_Heterogeneous(self.api_clients)
+        self.current_team_plan = None
 
-        if "ERROR:" in CLASSIFIER_SYSTEM_PROMPT:
-            raise FileNotFoundError(CLASSIFIER_SYSTEM_PROMPT)
+        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)..."
-        classification = await get_llm_response(
-            "Gemini", 
-            self.api_clients["Gemini"], 
-            CLASSIFIER_SYSTEM_PROMPT, 
-            problem
-        )
+        classification = await get_llm_response("Gemini", self.api_clients["Gemini"], CLASSIFIER_SYSTEM_PROMPT, problem)
         classification = classification.strip().replace("\"", "") 
         yield f"Diagnosis: {classification}"
 
+    @async_generator
+    async def _generate_and_evaluate(self, problem: str, classification: str, correction_prompt: Optional[str] = None):
+        solution_draft = ""
+        team_plan = {}
+        
+        if correction_prompt:
+            problem = f"{problem}\n\n{correction_prompt}"
+
+        default_persona = PERSONAS_DATA[config.DEFAULT_PERSONA_KEY]["description"]
+
+        if classification == "Direct_Procedure" or classification == "Holistic_Abstract_Reasoning":
+            if not correction_prompt:
+                await yield_("Deploying: Baseline Single Agent (Simplicity Hypothesis)...")
+            solution_draft = await self.single_agent.solve(problem, default_persona)
+            
+        elif classification == "Local_Geometric_Procedural":
+            if not correction_prompt:
+                await yield_("Deploying: Static Homogeneous Team (Expert Anomaly)...")
+            solution_draft = await self.homo_team.solve(problem, default_persona)
+            
+        elif classification == "Cognitive_Labyrinth":
+            if not correction_prompt:
+                await yield_("Deploying: Static Heterogeneous Team (Cognitive Diversity)...")
+                team_plan, calibration_errors = await self.calibrator.calibrate_team(problem)
+                if calibration_errors:
+                    await yield_("--- CALIBRATION WARNINGS ---")
+                    for err in calibration_errors: await yield_(err)
+                    await yield_("-----------------------------")
+                await yield_(f"Calibration complete. Best Team: {json.dumps({k: v['llm'] for k, v in team_plan.items()})}")
+                self.current_team_plan = team_plan
+            
+            # Reuse the calibrated team
+            solution_draft = await self.hetero_team.solve(problem, self.current_team_plan)
+        
+        else:
+            if not correction_prompt:
+                await yield_(f"Diagnosis '{classification}' is unknown. Defaulting to Single Agent.")
+            solution_draft = await self.single_agent.solve(problem, default_persona)
+
+        if "Error generating response" in solution_draft:
+            raise Exception(f"The specialist team failed to generate a solution. Error: {solution_draft}")
+
+        await yield_(f"Draft solution received: '{solution_draft[:60]}...'")
+
+        # --- EVALUATE ---
+        await yield_("Evaluating final draft (live)...")
+        v_fitness_json = await self.evaluator.evaluate(problem, solution_draft)
+        
+        # --- NEW: Robust Normalization of Evaluation Data ---
+        # This block fixes the "list object has no attribute get" error
+        normalized_fitness = {}
+        if isinstance(v_fitness_json, dict):
+            for k, v in v_fitness_json.items():
+                if isinstance(v, dict):
+                    normalized_fitness[k] = v
+                elif isinstance(v, list) and len(v) > 0 and isinstance(v[0], dict):
+                    # If the LLM wrapped the object in a list, unwrap it
+                    normalized_fitness[k] = v[0]
+                else:
+                    # Fallback for unexpected structure
+                    normalized_fitness[k] = {'score': 0, 'justification': str(v)}
+        else:
+             # Fallback if the whole thing isn't a dict
+             await yield_(f"Warning: Invalid JSON structure from Judge: {type(v_fitness_json)}")
+             normalized_fitness = {k: {'score': 0, 'justification': "Invalid JSON structure"} for k in ["Novelty", "Usefulness_Feasibility", "Flexibility", "Elaboration", "Cultural_Appropriateness"]}
+        
+        v_fitness_json = normalized_fitness
+        # ----------------------------------------------------
+
+        scores = {k: v.get('score', 0) for k, v in v_fitness_json.items()}
+        await yield_(f"Evaluation Score: {scores}")
+        
+        # Debug info if score is low
+        if scores.get('Novelty', 0) <= 1:
+             await yield_(f"⚠️ Low Score Detected. Reason: {v_fitness_json.get('Novelty', {}).get('justification', 'Unknown')}")
+        
+        return solution_draft, v_fitness_json, scores
+
     async def solve(self, problem: str) -> AsyncGenerator[str, None]:
         classification_generator = self._classify_problem(problem)
         classification = ""
@@ -147,60 +192,35 @@ class StrategicSelectorAgent:
             yield "Classifier failed. Defaulting to Single Agent."
             classification = "Direct_Procedure"
 
-        solution_draft = "" 
+        solution_draft, v_fitness_json, scores = "", {}, {}
         
         try:
-            default_persona = PERSONAS_DATA[config.DEFAULT_PERSONA_KEY]["description"]
-
-            if classification == "Direct_Procedure" or classification == "Holistic_Abstract_Reasoning":
-                yield "Deploying: Baseline Single Agent (Simplicity Hypothesis)..."
-                solution_draft = await self.single_agent.solve(problem, default_persona)
-                
-            elif classification == "Local_Geometric_Procedural":
-                yield "Deploying: Static Homogeneous Team (Expert Anomaly)..."
-                solution_draft = await self.homo_team.solve(problem, default_persona)
+            # --- MAIN LOOP (Self-Correction) ---
+            for i in range(2): 
+                current_problem = problem
+                if i > 0:
+                    yield f"--- (Loop {i}) Score is too low. Initiating Self-Correction... ---"
+                    correction_prompt_text = self.corrector.get_correction_plan(v_fitness_json)
+                    yield f"Diagnosis: {correction_prompt_text.splitlines()[3].strip()}"
+                    current_problem = f"{problem}\n\n{correction_prompt_text}"
                 
-            elif classification == "Cognitive_Labyrinth":
-                yield "Deploying: Static Heterogeneous Team (Cognitive Diversity)..."
+                loop_generator = self._generate_and_evaluate(current_problem, classification, None if i==0 else "Correcting...")
                 
-                # --- NEW: Capture errors from calibration ---
-                team_plan, calibration_errors = await self.calibrator.calibrate_team(problem)
+                async for status_update in loop_generator:
+                    yield status_update
                 
-                # --- 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)
-            
-            else:
-                yield f"Diagnosis '{classification}' is unknown. Defaulting to Single Agent."
-                solution_draft = await self.single_agent.solve(problem, default_persona)
-
-            if "Error generating response" in solution_draft:
-                raise Exception(f"The specialist team failed to generate a solution. Error: {solution_draft}")
-
-            yield f"Draft solution received: '{solution_draft[:60]}...'"
-
-            yield "Evaluating final draft (live)..."
-            v_fitness_json = await self.evaluator.evaluate(problem, solution_draft)
-            
-            scores = {k: v.get('score', 0) for k, v in v_fitness_json.items()}
-            yield f"Initial Score: {scores}"
-            # If the score is the default '1', show the error message hidden in the justification
-            if scores.get('Novelty', 0) == 1:
-                 yield f"⚠️ Low Score Detected. Reason: {v_fitness_json.get('Novelty', {}).get('justification', 'Unknown')}"
-            # -----------------------
-            
-            # --- This is where Milestone 5 will go ---
-            yield "Skipping self-correction for now..."
+                solution_draft, v_fitness_json, scores = await loop_generator.aclose() # Wait for return
+
+                # Check if we passed
+                if self.corrector.is_good_enough(scores):
+                    yield "--- Solution approved by self-corrector. ---"
+                    break
+                elif i == 1:
+                    yield "--- Max correction loops reached. Accepting best effort. ---"
             
+            # --- FINALIZE ---
             await asyncio.sleep(0.5)
-            yield "Milestone 4 (with error logging) Complete."
-            
+            yield "Milestone 5 Complete. Self-Correction loop is live."
             solution_draft_json_safe = json.dumps(solution_draft)
             yield f"FINAL: {{\"text\": {solution_draft_json_safe}, \"audio\": null}}"