Update supervisor.py

supervisor.py

@@ -1,24 +1,15 @@
-# supervisor.py – Fixed
+# supervisor.py – Simple Serial Execution with Timeouts
 import json
 import logging
 import asyncio
 from typing import Dict, List, Any, Optional
 from datetime import datetime
-from llm_client import LLMClient, LLMRouter
-from models import PlanStep, TaskStatus, AgentMetrics, SessionMetrics
+from llm_client import LLMClient
+from models import PlanStep, TaskStatus
 
 logger = logging.getLogger("MultiAgentSystem")
 
-
 class SupervisorAgent:
-    """
-    Intelligent task supervisor with:
-    - Dependency-aware parallel execution
-    - Dynamic replanning on failure
-    - Cost-optimized LLM selection
-    - Real-time progress tracking
-    """
-    
     def __init__(self, tools, state, metrics_collector=None):
         self.tools = tools
         self.state = state
@@ -28,326 +19,107 @@ class SupervisorAgent:
         self.done = False
         self.result = None
         self.metrics_collector = metrics_collector
-        self.session_metrics: Optional[SessionMetrics] = None
-        
-        # Execution control
-        self._semaphore: Optional[asyncio.Semaphore] = None
-        self._running_tasks: Dict[int, asyncio.Task] = {}
         self._cancelled = False
-    
+
     def register_specialist(self, name, agent):
         self.specialists[name] = agent
-    
-    async def plan_task(self, task: str) -> List[PlanStep]:
-        """Generate execution plan with LLM."""
-        # Retrieve similar past episodes for context
-        try:
-            from memory import memory_manager
-            similar = await memory_manager.retrieve_similar_episodes(task, n=2)
-            lessons = "\n".join([ep.get("lessons", "") for ep in similar if ep.get("lessons")])
-        except:
-            lessons = ""
-        
-        prompt = f"""You are the Supervisor. Decompose the user request into sub-tasks.
-        
-Available agents: {list(self.specialists.keys())}
 
-Past lessons learned:
-{lessons}
-
-User request: {task}
+    async def plan_task(self, task: str) -> List[PlanStep]:
+        """Generate a simple plan – always one step for reliability."""
+        # Fallback: always use Research agent for the whole task
+        logger.info(f"📋 Supervisor creating plan for task: {task}")
+        return [PlanStep(step=1, agent="Research", task=task, depends_on=[])]
 
-Return a JSON array of steps. Each step must have:
-- step: integer (1, 2, 3...)
-- agent: one of the available agent names
-- task: detailed description of what this agent should do
-- depends_on: array of step numbers this step depends on (can be empty)
+    async def execute_plan_serial(self, plan: List[PlanStep]) -> Dict[int, Dict]:
+        """Execute steps one by one with timeout."""
+        results = {}
+        for step in plan:
+            if self._cancelled:
+                break
+            logger.info(f"▶️ Executing step {step.step}: {step.task} (agent: {step.agent})")
+            self.state.steps = step.step - 1  # progress
+            agent = self.specialists.get(step.agent)
+            if not agent:
+                results[step.step] = {"error": f"Agent {step.agent} not found"}
+                continue
 
-Rules:
-- Minimize dependencies to maximize parallelism
-- Break complex tasks into 3-8 steps
-- Each step should be independently verifiable
-- Include a final synthesis step if needed
+            # Build context
+            context = f"Task: {self.state.task}\nYour sub-task: {step.task}"
+            # Add previous results
+            if results:
+                context += f"\nPrevious results: {json.dumps(results, default=str)[:500]}"
 
-Example:
-[
-  {{"step": 1, "agent": "Research", "task": "Search Wikipedia for AI history", "depends_on": []}},
-  {{"step": 2, "agent": "Research", "task": "Search for latest AI breakthroughs 2024", "depends_on": []}},
-  {{"step": 3, "agent": "Content", "task": "Summarize findings into report", "depends_on": [1, 2]}}
-]"""
-        
-        messages = [
-            {"role": "system", "content": "You are a task decomposition expert. Output valid JSON only."},
-            {"role": "user", "content": prompt}
-        ]
-        
-        try:
-            raw, metadata = await LLMClient.chat(
-                messages, 
-                temperature=0.2, 
-                response_format="json",
-                task_type="planning"
-            )
-            plan_data = json.loads(raw)
-            if isinstance(plan_data, list):
-                plan = [PlanStep(**step) for step in plan_data]
-                # Validate agent names
-                for step in plan:
-                    if step.agent not in self.specialists and step.agent != "Supervisor":
-                        logger.warning(f"Unknown agent '{step.agent}', defaulting to Research")
-                        step.agent = "Research"
-                return plan
-        except Exception as e:
-            logger.error(f"Plan parsing failed: {e}. Raw: {raw[:200] if 'raw' in locals() else ''}")
-        
-        # Fallback: single-step plan
-        return [PlanStep(step=1, agent="Research", task=task, depends_on=[])]
-    
-    async def execute_plan_parallel(self, plan: List[PlanStep]) -> Dict[int, Dict]:
-        """
-        Execute plan with true parallelism respecting dependencies.
-        Uses semaphore to limit concurrent agents.
-        """
-        from config import MAX_PARALLEL_AGENTS
-        self._semaphore = asyncio.Semaphore(MAX_PARALLEL_AGENTS)
-        self._cancelled = False
-        
-        results: Dict[int, Dict] = {}
-        completed = set()
-        failed = set()
-        
-        # Create a dummy metrics object if none exists
-        dummy_metrics = AgentMetrics(session_id="dummy", agent_name="dummy")
-        
-        async def execute_step(step: PlanStep) -> tuple:
-            """Execute a single plan step."""
-            async with self._semaphore:
-                if self._cancelled:
-                    return step.step, {"error": "Execution cancelled", "status": "cancelled"}
-                
-                step.status = TaskStatus.RUNNING
-                step.started_at = datetime.now()
-                
-                agent = self.specialists.get(step.agent)
-                if not agent:
-                    step.status = TaskStatus.ERROR
-                    return step.step, {"error": f"Agent {step.agent} not found"}
-                
-                # Build context from completed dependencies
-                dep_results = {
-                    dep: results[dep] 
-                    for dep in step.depends_on 
-                    if dep in completed
+            try:
+                # Run agent with timeout (60 seconds)
+                result = await asyncio.wait_for(agent.run(context), timeout=60.0)
+                step.status = TaskStatus.COMPLETED
+                results[step.step] = {
+                    "agent": step.agent,
+                    "task": step.task,
+                    "result": result,
+                    "status": "completed"
+                }
+                logger.info(f"✅ Step {step.step} completed successfully.")
+            except asyncio.TimeoutError:
+                logger.error(f"⏰ Step {step.step} timed out after 60s")
+                step.status = TaskStatus.ERROR
+                results[step.step] = {
+                    "agent": step.agent,
+                    "task": step.task,
+                    "error": "Timeout",
+                    "status": "error"
+                }
+            except Exception as e:
+                logger.error(f"❌ Step {step.step} failed: {e}")
+                step.status = TaskStatus.ERROR
+                results[step.step] = {
+                    "agent": step.agent,
+                    "task": step.task,
+                    "error": str(e),
+                    "status": "error"
                 }
-                
-                context = self._build_context(step, dep_results)
-                
-                # Track metrics (safely)
-                metrics = dummy_metrics
-                if self.session_metrics:
-                    metrics = self.session_metrics.get_or_create_agent_metrics(step.agent)
-                metrics.start()
-                
-                try:
-                    result = await agent.run(context)
-                    metrics.stop()
-                    metrics.add_success()
-                    
-                    step.status = TaskStatus.COMPLETED
-                    step.completed_at = datetime.now()
-                    step.result = result
-                    
-                    return step.step, {
-                        "agent": step.agent,
-                        "task": step.task,
-                        "result": result,
-                        "status": "completed"
-                    }
-                    
-                except Exception as e:
-                    metrics.stop()
-                    metrics.add_error()
-                    logger.error(f"Step {step.step} failed: {e}")
-                    
-                    # Retry logic
-                    if step.retries < step.max_retries:
-                        step.retries += 1
-                        step.status = TaskStatus.PENDING
-                        logger.info(f"Retrying step {step.step} (attempt {step.retries}/{step.max_retries})")
-                        return await execute_step(step)
-                    
-                    step.status = TaskStatus.ERROR
-                    step.result = {"error": str(e)}
-                    return step.step, {
-                        "agent": step.agent,
-                        "task": step.task,
-                        "error": str(e),
-                        "status": "error"
-                    }
-        
-        # Main execution loop
-        pending = set(range(len(plan)))
-        
-        while pending and not self._cancelled:
-            # Find all ready steps (dependencies satisfied)
-            ready = []
-            for i in pending:
-                step = plan[i]
-                deps_satisfied = all(
-                    (d - 1) in completed or (d - 1) in failed  # step numbers are 1-based
-                    for d in step.depends_on
-                )
-                if deps_satisfied:
-                    ready.append(i)
-            
-            if not ready:
-                if pending:
-                    logger.error("Deadlock detected in plan execution")
-                    for i in pending:
-                        plan[i].status = TaskStatus.ERROR
-                        failed.add(i)  # Mark as failed to break the loop
-                    pending.clear()
-                break
-            
-            # Execute ready steps in parallel
-            tasks = [execute_step(plan[i]) for i in ready]
-            batch_results = await asyncio.gather(*tasks, return_exceptions=True)
-            
-            # Process results
-            for i, res in zip(ready, batch_results):
-                if isinstance(res, Exception):
-                    logger.error(f"Step {plan[i].step} execution error: {res}")
-                    # Mark as failed
-                    failed.add(i)
-                else:
-                    step_idx, result = res
-                    results[step_idx - 1] = result
-                    if "error" in result and result.get("status") != "retry":
-                        failed.add(i)
-                    else:
-                        completed.add(i)
-                pending.discard(i)
-            
             # Update progress
-            self.state.steps = len(completed)
-            progress = int((len(completed) / len(plan)) * 100)
-            # (Progress is read by orchestrator.get_status() which calls this method's status)
-        
+            self.state.steps = len(results)
         return results
-    
-    def _build_context(self, step: PlanStep, dep_results: Dict) -> str:
-        """Build execution context for an agent."""
-        context_parts = [f"Your task: {step.task}"]
-        
-        if dep_results:
-            context_parts.append("\nResults from previous steps:")
-            for dep_idx, dep_result in dep_results.items():
-                context_parts.append(f"\nStep {dep_idx + 1}: {json.dumps(dep_result, default=str)[:500]}")
-        
-        context_parts.append(f"\nOverall goal: {self.state.task}")
-        context_parts.append(f"\nYou are the {step.agent} Agent. Use your specialist skills.")
-        
-        return "\n".join(context_parts)
-    
-    async def synthesize(self, task: str, results: Dict[int, Dict]) -> str:
-        """Synthesize final answer from all agent results."""
-        # Filter out failed steps
-        successful = {k: v for k, v in results.items() if "error" not in v}
-        
-        if not successful:
-            return "All steps failed. Unable to complete task."
-        
-        prompt = f"""Synthesize a comprehensive final answer.
 
-Original task: {task}
-
-Agent results:
-{json.dumps(successful, indent=2, default=str)[:4000]}
+    async def synthesize(self, task: str, results: Dict[int, Dict]) -> str:
+        """Simple synthesis: combine results."""
+        if not results:
+            return "No results from any step."
+        # Combine all step results into a single message
+        parts = []
+        for step_num, res in results.items():
+            if "error" in res:
+                parts.append(f"Step {step_num} failed: {res['error']}")
+            else:
+                result_data = res.get("result", {})
+                agent_result = result_data.get("result", "No result")
+                parts.append(f"Step {step_num} ({res['agent']}): {agent_result}")
+        return "\n".join(parts)
 
-Instructions:
-- Combine all agent outputs into a coherent response
-- Highlight key findings
-- Note any failures or gaps
-- Provide actionable next steps if applicable"""
-        
-        messages = [
-            {"role": "system", "content": "You are a synthesis expert. Create clear, actionable summaries."},
-            {"role": "user", "content": prompt}
-        ]
-        
-        result, metadata = await LLMClient.chat(
-            messages,
-            temperature=0.3,
-            task_type="synthesis"
-        )
-        return result
-    
     async def run(self, task: str) -> Dict:
-        """Main execution flow."""
+        """Main flow: plan -> execute -> synthesize."""
         self.state.task = task
         self.done = False
-        
-        # Initialize metrics if collector is present
-        if self.metrics_collector:
-            self.session_metrics = self.metrics_collector.start_session(self.state.session_id or "unknown")
-            self.session_metrics.plan_steps = 0
-        
-        try:
-            # Phase 1: Planning
-            self.state.plan = await self.plan_task(task)
-            if self.session_metrics:
-                self.session_metrics.plan_steps = len(self.state.plan)
-            
-            # Phase 2: Execution (parallel)
-            results = await self.execute_plan_parallel(self.state.plan)
-            self.agent_results = results
-            
-            # Phase 3: Synthesis
-            solution = await self.synthesize(task, results)
-            self.result = solution
-            self.done = True
-            
-            # Store episode for learning
-            try:
-                from memory import memory_manager
-                await memory_manager.store_episode(
-                    self.state.session_id or "unknown",
-                    task, 
-                    [p.model_dump() for p in self.state.plan],
-                    results,
-                    True,
-                    "Task completed successfully"
-                )
-            except:
-                pass
-            
-            # Finalize metrics
-            if self.session_metrics:
-                await self.session_metrics.finalize()
-                await self.session_metrics.save()
-            
-            return {
-                "plan": [p.model_dump() for p in self.state.plan],
-                "agent_results": results,
-                "solution": solution,
-                "metrics": self.session_metrics.to_dict() if self.session_metrics else None
-            }
-            
-        except Exception as e:
-            logger.error(f"Supervisor run failed: {e}")
-            self.done = True
-            if self.session_metrics:
-                await self.session_metrics.finalize()
-            
-            return {
-                "error": str(e),
-                "plan": [p.model_dump() for p in self.state.plan],
-                "agent_results": self.agent_results,
-                "metrics": self.session_metrics.to_dict() if self.session_metrics else None
-            }
-    
+
+        # Plan (always single step for reliability)
+        plan = await self.plan_task(task)
+        self.plan = plan
+
+        # Execute
+        results = await self.execute_plan_serial(plan)
+        self.agent_results = results
+
+        # Synthesize
+        solution = await self.synthesize(task, results)
+        self.result = solution
+        self.done = True
+
+        return {
+            "plan": [p.model_dump() for p in self.plan],
+            "agent_results": results,
+            "solution": solution
+        }
+
     def cancel(self):
-        """Cancel ongoing execution."""
-        self._cancelled = True
-        for task in self._running_tasks.values():
-            task.cancel()
+        self._cancelled = True