Update agent_logic.py

agent_logic.py

@@ -1,4 +1,4 @@
-# agent_logic.py (Milestone 5 - FINAL & ROBUST)
+# agent_logic.py (FINAL, Robust Version)
 import asyncio
 from typing import AsyncGenerator, Dict, Optional
 import json
@@ -6,13 +6,14 @@ import os
 import google.generativeai as genai
 from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
+import re # <-- Added: For score parsing fix
 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_
 
+# (Configuration and Manager Prompts are loaded here)
 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"])
@@ -79,8 +80,7 @@ class StrategicSelectorAgent:
             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)
+                self.api_clients["SambaNova"] = AsyncOpenAI(api_key=api_keys["sambanova"], base_url=os.getenv("SAMBANOVA_BASE_URL", "https://api.sambanova.ai/v1"))
             except Exception as e: print(f"Warning: SambaNova init failed: {e}")
         
         if not self.api_clients["Gemini"]:
@@ -102,84 +102,6 @@ class StrategicSelectorAgent:
         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 = ""
@@ -192,11 +114,14 @@ class StrategicSelectorAgent:
             yield "Classifier failed. Defaulting to Single Agent."
             classification = "Direct_Procedure"
 
-        solution_draft, v_fitness_json, scores = "", {}, {}
+        solution_draft = "" 
+        v_fitness_json = {}
+        scores = {}
         
         try:
             # --- 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... ---"
@@ -204,12 +129,87 @@ class StrategicSelectorAgent:
                     yield f"Diagnosis: {correction_prompt_text.splitlines()[3].strip()}"
                     current_problem = f"{problem}\n\n{correction_prompt_text}"
                 
-                loop_generator = self._generate_and_evaluate(current_problem, classification, None if i==0 else "Correcting...")
+                # --- DEPLOY ---
+                default_persona = PERSONAS_DATA[config.DEFAULT_PERSONA_KEY]["description"]
+
+                if classification == "Direct_Procedure" or classification == "Holistic_Abstract_Reasoning":
+                    if i == 0: yield "Deploying: Baseline Single Agent (Simplicity Hypothesis)..."
+                    solution_draft = await self.single_agent.solve(current_problem, default_persona)
+                    
+                elif classification == "Local_Geometric_Procedural":
+                    if i == 0: yield "Deploying: Static Homogeneous Team (Expert Anomaly)..."
+                    solution_draft = await self.homo_team.solve(current_problem, default_persona)
+                    
+                elif classification == "Cognitive_Labyrinth":
+                    if i == 0:
+                        yield "Deploying: Static Heterogeneous Team (Cognitive Diversity)..."
+                        team_plan, calibration_errors = await self.calibrator.calibrate_team(current_problem)
+                        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()})}"
+                        self.current_team_plan = team_plan
+                    
+                    solution_draft = await self.hetero_team.solve(current_problem, self.current_team_plan)
                 
-                async for status_update in loop_generator:
-                    yield status_update
+                else:
+                    if i == 0: yield f"Diagnosis '{classification}' is unknown. Defaulting to Single Agent."
+                    solution_draft = await self.single_agent.solve(current_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]}...'"
+
+                # --- EVALUATE ---
+                yield "Evaluating draft (live)..."
+                v_fitness_json = await self.evaluator.evaluate(current_problem, solution_draft)
+                
+                # --- Robust Normalization of Evaluation Data ---
+                normalized_fitness = {}
+                if isinstance(v_fitness_json, dict):
+                    for k, v in v_fitness_json.items():
+                        if isinstance(v, dict):
+                            # Standard format: {"score": 4, "justification": "..."}
+                            score_value = v.get('score')
+                            justification_value = v.get('justification', str(v))
+                        elif isinstance(v, list) and len(v) > 0 and isinstance(v[0], dict):
+                            # Handles list structure: [{"score": 4, "justification": "..."}]
+                            score_value = v[0].get('score')
+                            justification_value = v[0].get('justification', str(v[0]))
+                        else:
+                            # Fallback for unexpected structure
+                            score_value = 0
+                            justification_value = str(v)
+                        
+                        # FIX: Extract the integer score from the string (e.g., "4/5" -> 4)
+                        if isinstance(score_value, str):
+                            try:
+                                score_value = int(re.search(r'\d+', score_value).group())
+                            except:
+                                score_value = 0
+                        
+                        # Ensure score is an integer
+                        try:
+                            score_value = int(score_value)
+                        except (ValueError, TypeError):
+                            score_value = 0
+
+                        normalized_fitness[k] = {'score': score_value, 'justification': justification_value}
+
+                else:
+                     # Fallback if the whole thing isn't a dict
+                     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()}
+                yield f"Evaluation Score: {scores}"
                 
-                solution_draft, v_fitness_json, scores = await loop_generator.aclose() # Wait for return
+                if scores.get('Novelty', 0) <= 1:
+                     yield f"⚠️ Low Score Detected. Reason: {v_fitness_json.get('Novelty', {}).get('justification', 'Unknown')}"
 
                 # Check if we passed
                 if self.corrector.is_good_enough(scores):