agent_logic.py

# agent_logic.py (Updated for Milestone 3 - FIX)
import asyncio
from typing import AsyncGenerator

from mcp_servers import AgentCalibrator, BusinessSolutionEvaluator
# --- 1. Import the persona data dictionary ---
from personas import PERSONAS_DATA 


# --- 2. Define the Specialist Agent Architectures ---

class Baseline_Single_Agent:
    """Deploys a single agent (from Paper 2)"""
    async def solve(self, problem: str, persona_prompt: str):
        print(f"--- (Specialist Team: Single Agent) solving... ---")
        await asyncio.sleep(1) # Simulate LLM call
        # Placeholder: In Milestone 4, this will call Gemini
        return f"Solution from Single Agent ({persona_prompt.splitlines()[1][:20]}...)"

class Baseline_Static_Homogeneous:
    """Deploys a team of 3 agents with the same persona (from Paper 2)"""
    async def solve(self, problem: str, persona_prompt: str):
        print(f"--- (Specialist Team: Homogeneous) solving... ---")
        await asyncio.sleep(1.5) # Simulate 3 LLM calls + manager call
        # Placeholder: In Milestone 4, this will call 3 LLMs and a manager
        return f"Solution from Homogeneous Team ({persona_prompt.splitlines()[1][:20]}...)"

class Baseline_Static_Heterogeneous:
    """Deploys a diverse team with calibrated roles (from Paper 1 & 2)"""
    async def solve(self, problem: str, team_plan: dict):
        print(f"--- (Specialist Team: Heterogeneous) solving... ---")
        # team_plan = {
        #   "Plant": {"persona": "Culture_5", "llm": "Anthropic"},
        #   "Implementer": {"persona": "Culture_Expert", "llm": "Gemini"},
        #   "Monitor": {"persona": "Culture_11", "llm": "SambaNova"}
        # }
        print(f"Deploying team with plan: {team_plan}")
        await asyncio.sleep(1.5) # Simulate 3 LLM calls + manager call
        # Placeholder: In Milestone 4, this will call 3 different LLMs
        return f"Solution from Heterogeneous Team (Plant: {team_plan['Plant']['llm']})"


class StrategicSelectorAgent:
    """
    This is MudabbirAI, the orchestrator.
    It diagnoses the problem and deploys the correct agent team.
    """
    def __init__(self):
        self.evaluator = BusinessSolutionEvaluator()
        self.calibrator = AgentCalibrator(self.evaluator)
        
        # --- 3. Instantiate the new specialist teams ---
        self.single_agent = Baseline_Single_Agent()
        self.homo_team = Baseline_Static_Homogeneous()
        self.hetero_team = Baseline_Static_Heterogeneous()


    async def _classify_problem(self, problem: str) -> AsyncGenerator[str, None]:
        """
        Runs "OODA Reconnaissance" (Paper 2) to classify the problem.
        (Hard-coded for now)
        """
        yield "Classifying problem archetype..."
        await asyncio.sleep(1) # Simulate LLM call
        classification = "Cognitive_Labyrinth" 
        yield f"Diagnosis: {classification}"


    async def solve(self, problem: str) -> AsyncGenerator[str, None]:
        """
        Main generator function that orchestrates the entire solution.
        """
        
        # --- 1. CLASSIFY (from Paper 2) ---
        classification_generator = self._classify_problem(problem)
        classification = ""
        async for status_update in classification_generator:
            yield status_update
            if "Diagnosis: " in status_update:
                classification = status_update.split(": ")[-1]

        # --- 4. Update DEPLOY logic to call the teams ---
        solution_draft = "" 
        
        if classification == "Direct_Procedure" or classification == "Holistic_Abstract_Reasoning":
            yield "Deploying: Baseline Single Agent (Simplicity Hypothesis)..."
            # --- FIX: Access dictionary directly ---
            persona = PERSONAS_DATA["Culture_Expert"]["description"]
            solution_draft = await self.single_agent.solve(problem, persona)
            
        elif classification == "Local_Geometric_Procedural":
            yield "Deploying: Static Homogeneous Team (Expert Anomaly)..."
            # --- FIX: Access dictionary directly ---
            persona = PERSONAS_DATA["Culture_Expert"]["description"] # The "Implementer"
            solution_draft = await self.homo_team.solve(problem, persona)
            
        elif classification == "Cognitive_Labyrinth":
            yield "Deploying: Static Heterogeneous Team (Cognitive Diversity)..."
            
            # CALIBRATE step
            team_plan = await self.calibrator.calibrate_team(problem)
            yield f"Calibration complete. Best Plant: {team_plan['Plant']['llm']}"
            
            # DEPLOY calibrated team
            solution_draft = await self.hetero_team.solve(problem, team_plan)
        
        else:
            yield f"Error: Unknown problem classification '{classification}'."
            return

        yield f"Draft solution received: '{solution_draft[:50]}...'"

        # --- EVALUATE ---
        yield "Evaluating final draft..."
        v_fitness = await self.evaluator.evaluate(solution_draft)
        yield f"Initial Score: {v_fitness['Novelty']} Novelty, {v_fitness['Usefulness_Feasibility']} Feasibility"
        
        # --- (SKIPPED) SELF-CORRECT ---
        yield "Skipping self-correction for Milestone 3..."
        
        # --- FINALIZE ---
        await asyncio.sleep(0.5) # Simulate work
        yield "Milestone 3 Complete. Specialist teams are integrated."
        yield f"FINAL: {{\"text\": \"{solution_draft}\", \"audio\": null}}"