import time from typing import Dict, List, Optional, Union from app.agent.base import BaseAgent from app.flow.base import BaseFlow from app.logger import logger from app.schema import AgentState from app.tool import PlanningTool, ToolCollection class PlanningFlow(BaseFlow): """A flow that manages planning and execution of tasks using agents.""" def __init__( self, agents: Union[BaseAgent, List[BaseAgent], Dict[str, BaseAgent]], **kwargs ): # Initialize planning tool first self.planning_tool = self._initialize_planning_tool(kwargs.get("tools")) # If tools were provided, ensure planning tool is included tools = kwargs.get("tools") if tools: planning_tool_exists = any( isinstance(tool, PlanningTool) for tool in tools.tools ) if not planning_tool_exists: tools.add_tool(self.planning_tool) else: # Create a new tool collection with at least the planning tool tools = ToolCollection(self.planning_tool) kwargs["tools"] = tools super().__init__(agents, **kwargs) # Define agent roles self.planner_key = kwargs.get("planner", self.primary_agent_key) self.executor_keys = kwargs.get("executors", list(self.agents.keys())) # Planning state tracking self.active_plan_id = kwargs.get("plan_id", f"plan_{int(time.time())}") self.current_step_index = None # Ensure the planning tool has been initialized properly if not hasattr(self.planning_tool, "_plans"): self.planning_tool._plans = {} def _initialize_planning_tool( self, tools: Optional[ToolCollection] ) -> PlanningTool: """Initialize planning tool, reusing existing one if available""" if tools: for tool in tools.tools: if isinstance(tool, PlanningTool): return tool return PlanningTool() @property def planner(self) -> Optional[BaseAgent]: """Get the planning agent""" return ( self.agents.get(self.planner_key) if self.planner_key else self.primary_agent ) def get_executor(self, step_type: Optional[str] = None) -> BaseAgent: """ Get an appropriate executor agent for the current step. Can be extended to select agents based on step type/requirements. """ # If step type is provided and matches an agent key, use that agent if step_type and step_type in self.agents: return self.agents[step_type] # Otherwise use the first available executor or fall back to primary agent for key in self.executor_keys: if key in self.agents: return self.agents[key] # Fallback to primary agent return self.primary_agent async def execute(self, input_text: str) -> str: """Execute the planning flow with agents.""" try: if not self.primary_agent: raise ValueError("No primary agent available") # Create initial plan if input provided if input_text: await self._create_initial_plan(input_text) # Verify plan was created successfully if self.active_plan_id not in self.planning_tool._plans: logger.error( f"Plan creation failed. Plan ID {self.active_plan_id} not found in planning tool." ) return f"Failed to create plan for: {input_text}" result = "" while True: # Get current step to execute self.current_step_index, step_info = await self._get_current_step_info() # Exit if no more steps or plan completed if self.current_step_index is None: result += await self._finalize_plan() break # Execute current step with appropriate agent step_type = step_info.get("type") if step_info else None executor = self.get_executor(step_type) step_result = await self._execute_step(executor, step_info) result += step_result + "\n" # Check if agent wants to terminate if hasattr(executor, "state") and executor.state == AgentState.FINISHED: break return result except Exception as e: logger.error(f"Error in PlanningFlow: {str(e)}") return f"Execution failed: {str(e)}" async def _create_initial_plan(self, request: str) -> None: """Create an initial plan based on the request using an appropriate agent.""" logger.info(f"Creating initial plan with ID: {self.active_plan_id}") agent = self.planner if self.planner else self.primary_agent # First, directly create an empty plan to ensure the plan ID exists self.planning_tool._plans[self.active_plan_id] = { "title": f"Plan for: {request[:50]}{'...' if len(request) > 50 else ''}", "description": f"Auto-generated plan for request: {request}", "steps": [], "step_status": {}, "created_at": time.time(), "updated_at": time.time(), } # Use agent.run to create the plan plan_prompt = f""" I need you to create a detailed plan to accomplish this task: {request} Please create a plan with ID {self.active_plan_id} using the planning tool. The plan should include all necessary steps to complete the task. """ try: plan_result = await agent.run(plan_prompt) logger.info(f"Plan creation result: {plan_result[:200]}...") # Verify the plan was created if ( self.active_plan_id not in self.planning_tool._plans or not self.planning_tool._plans[self.active_plan_id].get("steps") ): logger.warning( "Plan may not have been created properly. Creating default plan." ) await self._create_default_plan(request) except Exception as e: logger.error(f"Error creating plan: {e}") await self._create_default_plan(request) async def _create_default_plan(self, request: str) -> None: """Create a default plan if the agent fails to create one.""" try: # Try using the planning tool directly await self.planning_tool.execute( command="create", plan_id=self.active_plan_id, title=f"Plan for: {request[:50]}{'...' if len(request) > 50 else ''}", description=f"Auto-generated plan for request: {request}", steps=["Analyze request", "Execute task", "Verify results"], ) except Exception as e: logger.error(f"Failed to create default plan with planning tool: {e}") # Create plan directly in the planning tool's storage self.planning_tool._plans[self.active_plan_id] = { "title": f"Emergency Plan for: {request[:50]}{'...' if len(request) > 50 else ''}", "description": f"Emergency auto-generated plan for request: {request}", "steps": ["Analyze request", "Execute task", "Verify results"], "step_status": { "0": "not_started", "1": "not_started", "2": "not_started", }, "created_at": time.time(), "updated_at": time.time(), } logger.info(f"Created default plan with ID: {self.active_plan_id}") async def _get_current_step_info(self) -> tuple[Optional[int], Optional[dict]]: """ Parse the current plan to identify the first non-completed step's index and info. Returns (None, None) if no active step is found. """ if ( not self.active_plan_id or self.active_plan_id not in self.planning_tool._plans ): logger.error(f"Plan with ID {self.active_plan_id} not found") return None, None try: # Direct access to step status from planning tool storage plan_data = self.planning_tool._plans[self.active_plan_id] steps = plan_data.get("steps", []) step_status = plan_data.get("step_status", {}) # Find first non-completed step for i, step in enumerate(steps): status = step_status.get(str(i), "not_started") if status in ["not_started", "in_progress"]: # Extract step type/category if available step_info = {"text": step} # Try to extract step type from the text (e.g., [SEARCH] or [CODE]) import re type_match = re.search(r"\[([A-Z_]+)\]", step) if type_match: step_info["type"] = type_match.group(1).lower() # Mark current step as in_progress try: await self.planning_tool.execute( command="mark_step", plan_id=self.active_plan_id, step_index=i, step_status="in_progress", ) except Exception as e: logger.warning(f"Error marking step as in_progress: {e}") # Update step status directly step_status[str(i)] = "in_progress" plan_data["step_status"] = step_status plan_data["updated_at"] = time.time() return i, step_info return None, None # No active step found except Exception as e: logger.warning(f"Error finding current step index: {e}") return None, None async def _execute_step(self, executor: BaseAgent, step_info: dict) -> str: """Execute the current step with the specified agent using agent.run().""" # Prepare context for the agent with current plan status plan_status = await self._get_plan_text() step_text = step_info.get("text", f"Step {self.current_step_index}") # Create a prompt for the agent to execute the current step step_prompt = f""" CURRENT PLAN STATUS: {plan_status} YOUR CURRENT TASK: You are now working on step {self.current_step_index}: "{step_text}" Please execute this step using the appropriate tools. When you're done, provide a summary of what you accomplished. """ # Use agent.run() to execute the step try: step_result = await executor.run(step_prompt) # Mark the step as completed after successful execution await self._mark_step_completed() return step_result except Exception as e: logger.error(f"Error executing step {self.current_step_index}: {e}") return f"Error executing step {self.current_step_index}: {str(e)}" async def _mark_step_completed(self) -> None: """Mark the current step as completed.""" if self.current_step_index is None: return try: # Mark the step as completed await self.planning_tool.execute( command="mark_step", plan_id=self.active_plan_id, step_index=self.current_step_index, step_status="completed", ) logger.info( f"Marked step {self.current_step_index} as completed in plan {self.active_plan_id}" ) except Exception as e: logger.warning(f"Failed to update plan status: {e}") # Update step status directly in planning tool storage if self.active_plan_id in self.planning_tool._plans: plan_data = self.planning_tool._plans[self.active_plan_id] step_status = plan_data.get("step_status", {}) step_status[str(self.current_step_index)] = "completed" plan_data["step_status"] = step_status plan_data["updated_at"] = time.time() async def _get_plan_text(self) -> str: """Get the current plan as formatted text.""" try: result = await self.planning_tool.execute( command="get", plan_id=self.active_plan_id ) return result.output if hasattr(result, "output") else str(result) except Exception as e: logger.error(f"Error getting plan: {e}") return self._generate_plan_text_from_storage() def _generate_plan_text_from_storage(self) -> str: """Generate plan text directly from storage if the planning tool fails.""" try: if self.active_plan_id not in self.planning_tool._plans: return f"Error: Plan with ID {self.active_plan_id} not found" plan_data = self.planning_tool._plans[self.active_plan_id] title = plan_data.get("title", "Untitled Plan") description = plan_data.get("description", "") steps = plan_data.get("steps", []) step_status = plan_data.get("step_status", {}) # Count steps by status status_counts = { "completed": 0, "in_progress": 0, "blocked": 0, "not_started": 0, } for status in step_status.values(): if status in status_counts: status_counts[status] += 1 completed = status_counts["completed"] total = len(steps) progress = (completed / total) * 100 if total > 0 else 0 plan_text = f"Plan: {title} (ID: {self.active_plan_id})\n" plan_text += "=" * len(plan_text) + "\n\n" plan_text += f"{description}\n\n" if description else "" plan_text += ( f"Progress: {completed}/{total} steps completed ({progress:.1f}%)\n" ) plan_text += f"Status: {status_counts['completed']} completed, {status_counts['in_progress']} in progress, " plan_text += f"{status_counts['blocked']} blocked, {status_counts['not_started']} not started\n\n" plan_text += "Steps:\n" for i, step in enumerate(steps): status = step_status.get(str(i), "not_started") if status == "completed": status_mark = "[✓]" elif status == "in_progress": status_mark = "[→]" elif status == "blocked": status_mark = "[!]" else: # not_started status_mark = "[ ]" plan_text += f"{i}. {status_mark} {step}\n" return plan_text except Exception as e: logger.error(f"Error generating plan text from storage: {e}") return f"Error: Unable to retrieve plan with ID {self.active_plan_id}" async def _get_plan(self) -> dict: """Get the current plan as a dictionary.""" if ( not self.active_plan_id or self.active_plan_id not in self.planning_tool._plans ): return {} return self.planning_tool._plans[self.active_plan_id] async def _finalize_plan(self) -> str: """Finalize the plan and provide a summary using an appropriate agent.""" agent = self.planner if self.planner else self.primary_agent plan_text = await self._get_plan_text() # Create a summary prompt summary_prompt = f""" The plan has been completed. Here is the final plan status: {plan_text} Please provide a summary of what was accomplished and any final thoughts. """ # Use agent.run() to generate the summary try: summary = await agent.run(summary_prompt) return f"Plan completed:\n\n{summary}" except Exception as e: logger.error(f"Error finalizing plan: {e}") return "Plan completed. Error generating summary."