fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl

mcp_agent/agents/workflow/orchestrator_agent.py

@@ -0,0 +1,591 @@
+"""
+OrchestratorAgent implementation using the BaseAgent adapter pattern.
+
+This workflow provides an implementation that manages complex tasks by
+dynamically planning, delegating to specialized agents, and synthesizing results.
+"""
+
+from typing import Any, Dict, List, Literal, Optional, Type
+
+from mcp.types import TextContent
+
+from mcp_agent.agents.agent import Agent
+from mcp_agent.agents.base_agent import BaseAgent
+from mcp_agent.agents.workflow.orchestrator_models import (
+    NextStep,
+    Plan,
+    PlanResult,
+    Step,
+    TaskWithResult,
+    format_plan_result,
+    format_step_result_text,
+)
+from mcp_agent.agents.workflow.orchestrator_prompts import (
+    FULL_PLAN_PROMPT_TEMPLATE,
+    ITERATIVE_PLAN_PROMPT_TEMPLATE,
+    SYNTHESIZE_INCOMPLETE_PLAN_TEMPLATE,
+    SYNTHESIZE_PLAN_PROMPT_TEMPLATE,
+    TASK_PROMPT_TEMPLATE,
+)
+from mcp_agent.core.agent_types import AgentConfig
+from mcp_agent.core.exceptions import AgentConfigError
+from mcp_agent.core.request_params import RequestParams
+from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.interfaces import ModelT
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+
+logger = get_logger(__name__)
+
+
+class OrchestratorAgent(BaseAgent):
+    """
+    An agent that implements the orchestrator workflow pattern.
+
+    Dynamically creates execution plans and delegates tasks
+    to specialized worker agents, synthesizing their results into a cohesive output.
+    Supports both full planning and iterative planning modes.
+    """
+
+    def __init__(
+        self,
+        config: AgentConfig,
+        agents: List[Agent],
+        plan_type: Literal["full", "iterative"] = "full",
+        context: Optional[Any] = None,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize an OrchestratorAgent.
+
+        Args:
+            config: Agent configuration or name
+            agents: List of specialized worker agents available for task execution
+            plan_type: Planning mode ("full" or "iterative")
+            context: Optional context object
+            **kwargs: Additional keyword arguments to pass to BaseAgent
+        """
+        super().__init__(config, context=context, **kwargs)
+
+        if not agents:
+            raise AgentConfigError("At least one worker agent must be provided")
+
+        self.plan_type = plan_type
+
+        # Store agents by name for easier lookup
+        self.agents: Dict[str, Agent] = {}
+        for agent in agents:
+            agent_name = agent.name
+            self.logger.info(f"Adding agent '{agent_name}' to orchestrator")
+            self.agents[agent_name] = agent
+
+        # For tracking state during execution
+        self.plan_result: Optional[PlanResult] = None
+
+    async def generate(
+        self,
+        multipart_messages: List[PromptMessageMultipart],
+        request_params: Optional[RequestParams] = None,
+    ) -> PromptMessageMultipart:
+        """
+        Execute an orchestrated plan to process the input.
+
+        Args:
+            multipart_messages: Messages to process
+            request_params: Optional request parameters
+
+        Returns:
+            The final synthesized response from the orchestration
+        """
+        # Extract user request
+        objective = multipart_messages[-1].all_text() if multipart_messages else ""
+
+        # Initialize execution parameters
+        params = self._merge_request_params(request_params)
+
+        # Execute the plan
+        plan_result = await self._execute_plan(objective, params)
+        self.plan_result = plan_result
+
+        # Return the result
+        return PromptMessageMultipart(
+            role="assistant",
+            content=[TextContent(type="text", text=plan_result.result or "No result available")],
+        )
+
+    async def structured(
+        self,
+        prompt: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: Optional[RequestParams] = None,
+    ) -> Optional[ModelT]:
+        """
+        Execute an orchestration plan and parse the result into a structured format.
+
+        Args:
+            prompt: List of messages to process
+            model: Pydantic model to parse the response into
+            request_params: Optional request parameters
+
+        Returns:
+            The parsed final response, or None if parsing fails
+        """
+        # Generate orchestration result
+        response = await self.generate(prompt, request_params)
+
+        # Try to parse the response into the specified model
+        try:
+            result_text = response.all_text()
+            prompt_message = PromptMessageMultipart(
+                role="user", content=[TextContent(type="text", text=result_text)]
+            )
+
+            # Use the LLM's structured parsing capability
+            return await self._llm.structured([prompt_message], model, request_params)
+        except Exception as e:
+            self.logger.warning(f"Failed to parse orchestration result: {str(e)}")
+            return None
+
+    async def initialize(self) -> None:
+        """Initialize the orchestrator agent and worker agents."""
+        await super().initialize()
+
+        # Initialize all worker agents if not already initialized
+        for agent_name, agent in self.agents.items():
+            if not getattr(agent, "initialized", False):
+                self.logger.debug(f"Initializing agent: {agent_name}")
+                await agent.initialize()
+
+        self.initialized = True
+
+    async def shutdown(self) -> None:
+        """Shutdown the orchestrator agent and worker agents."""
+        await super().shutdown()
+
+        # Shutdown all worker agents
+        for agent_name, agent in self.agents.items():
+            try:
+                await agent.shutdown()
+            except Exception as e:
+                self.logger.warning(f"Error shutting down agent {agent_name}: {str(e)}")
+
+    async def _execute_plan(self, objective: str, request_params: RequestParams) -> PlanResult:
+        """
+        Execute a plan to achieve the given objective.
+
+        Args:
+            objective: The objective to achieve
+            request_params: Request parameters for execution
+
+        Returns:
+            PlanResult containing execution results and final output
+        """
+        iterations = 0
+        total_steps_executed = 0
+        max_iterations = request_params.max_iterations
+        max_steps = getattr(request_params, "max_steps", max_iterations * 5)
+
+        # Initialize plan result
+        plan_result = PlanResult(objective=objective, step_results=[])
+        plan_result.max_iterations_reached = False
+
+        while iterations < max_iterations:
+            # Generate plan based on planning mode
+            if self.plan_type == "iterative":
+                next_step = await self._get_next_step(objective, plan_result, request_params)
+                if next_step is None:
+                    self.logger.error("Failed to generate next step, ending iteration early")
+                    plan_result.max_iterations_reached = True
+                    break
+
+                logger.debug(f"Iteration {iterations}: Iterative plan:", data=next_step)
+                plan = Plan(steps=[next_step], is_complete=next_step.is_complete)
+            elif self.plan_type == "full":
+                plan = await self._get_full_plan(objective, plan_result, request_params)
+                if plan is None:
+                    self.logger.error("Failed to generate full plan, ending iteration early")
+                    plan_result.max_iterations_reached = True
+                    break
+
+                logger.debug(f"Iteration {iterations}: Full Plan:", data=plan)
+            else:
+                raise ValueError(f"Invalid plan type: {self.plan_type}")
+
+            # Validate agent names early
+            self._validate_agent_names(plan)
+
+            # Store plan in result
+            plan_result.plan = plan
+
+            # Execute the steps in the plan
+            for step in plan.steps:
+                # Check if we've hit the step limit
+                if total_steps_executed >= max_steps:
+                    self.logger.warning(
+                        f"Reached maximum step limit ({max_steps}) without completing objective"
+                    )
+                    plan_result.max_steps_reached = True
+                    break
+
+                # Execute the step and collect results
+                step_result = await self._execute_step(step, plan_result, request_params)
+
+                plan_result.add_step_result(step_result)
+                total_steps_executed += 1
+
+            # Check if we need to break due to hitting max steps
+            if getattr(plan_result, "max_steps_reached", False):
+                break
+
+            # If the plan is marked complete, finalize the result
+            if plan.is_complete:
+                plan_result.is_complete = True
+                break
+
+            # Increment iteration counter
+            iterations += 1
+
+        # Generate final result based on execution status
+        if iterations >= max_iterations and not plan_result.is_complete:
+            self.logger.warning(f"Failed to complete in {max_iterations} iterations")
+            plan_result.max_iterations_reached = True
+
+            # Use incomplete plan template
+            synthesis_prompt = SYNTHESIZE_INCOMPLETE_PLAN_TEMPLATE.format(
+                plan_result=format_plan_result(plan_result), max_iterations=max_iterations
+            )
+        else:
+            # Either plan is complete or we had other limits
+            if not plan_result.is_complete:
+                plan_result.is_complete = True
+
+            # Use standard template
+            synthesis_prompt = SYNTHESIZE_PLAN_PROMPT_TEMPLATE.format(
+                plan_result=format_plan_result(plan_result)
+            )
+
+        # Generate final synthesis
+        plan_result.result = await self._planner_generate_str(
+            synthesis_prompt, request_params.model_copy(update={"max_iterations": 1})
+        )
+
+        return plan_result
+
+    async def _execute_step(
+        self, step: Step, previous_result: PlanResult, request_params: RequestParams
+    ) -> Any:
+        """
+        Execute a single step from the plan.
+
+        Args:
+            step: The step to execute
+            previous_result: Results of the plan execution so far
+            request_params: Request parameters
+
+        Returns:
+            Result of executing the step
+        """
+        from mcp_agent.agents.workflow.orchestrator_models import StepResult
+
+        # Initialize step result
+        step_result = StepResult(step=step, task_results=[])
+
+        # Format context for tasks
+        context = format_plan_result(previous_result)
+
+        # Execute all tasks in parallel
+        futures = []
+        error_tasks = []
+
+        for task in step.tasks:
+            # Check agent exists
+            agent = self.agents.get(task.agent)
+            if not agent:
+                self.logger.error(
+                    f"No agent found matching '{task.agent}'. Available agents: {list(self.agents.keys())}"
+                )
+                error_tasks.append(
+                    (
+                        task,
+                        f"Error: Agent '{task.agent}' not found. Available agents: {', '.join(self.agents.keys())}",
+                    )
+                )
+                continue
+
+            # Prepare task prompt
+            task_description = TASK_PROMPT_TEMPLATE.format(
+                objective=previous_result.objective, task=task.description, context=context
+            )
+
+            # Queue task for execution
+            futures.append(
+                (
+                    task,
+                    agent.generate(
+                        [
+                            PromptMessageMultipart(
+                                role="user",
+                                content=[TextContent(type="text", text=task_description)],
+                            )
+                        ]
+                    ),
+                )
+            )
+
+        # Wait for all tasks
+        task_results = []
+        for future in futures:
+            task, future_obj = future
+            try:
+                result = await future_obj
+                result_text = result.all_text()
+
+                # Create task result
+                task_model = task.model_dump()
+                task_result = TaskWithResult(
+                    description=task_model["description"],
+                    agent=task_model["agent"],
+                    result=result_text,
+                )
+                task_results.append(task_result)
+            except Exception as e:
+                self.logger.error(f"Error executing task: {str(e)}")
+                # Add error result
+                task_model = task.model_dump()
+                task_results.append(
+                    TaskWithResult(
+                        description=task_model["description"],
+                        agent=task_model["agent"],
+                        result=f"ERROR: {str(e)}",
+                    )
+                )
+
+        # Add all task results to step result
+        for task_result in task_results:
+            step_result.add_task_result(task_result)
+
+        # Add error task results
+        for task, error_message in error_tasks:
+            task_model = task.model_dump()
+            step_result.add_task_result(
+                TaskWithResult(
+                    description=task_model["description"],
+                    agent=task_model["agent"],
+                    result=f"ERROR: {error_message}",
+                )
+            )
+
+        # Format step result
+        step_result.result = format_step_result_text(step_result)
+        return step_result
+
+    async def _get_full_plan(
+        self, objective: str, plan_result: PlanResult, request_params: RequestParams
+    ) -> Optional[Plan]:
+        """
+        Generate a full plan with all steps.
+
+        Args:
+            objective: The objective to achieve
+            plan_result: Current plan execution state
+            request_params: Request parameters
+
+        Returns:
+            Complete Plan with all steps, or None if parsing fails
+        """
+        # Format agent information for the prompt
+        agent_formats = []
+        for agent_name in self.agents.keys():
+            formatted = self._format_agent_info(agent_name)
+            agent_formats.append(formatted)
+
+        agents = "\n".join(agent_formats)
+
+        # Determine plan status
+        if plan_result.is_complete:
+            plan_status = "Plan Status: Complete"
+        elif plan_result.step_results:
+            plan_status = "Plan Status: In Progress"
+        else:
+            plan_status = "Plan Status: Not Started"
+
+        # Calculate iteration information
+        max_iterations = request_params.max_iterations
+        current_iteration = len(plan_result.step_results)
+        current_iteration = min(current_iteration, max_iterations - 1)
+        iterations_remaining = max(0, max_iterations - current_iteration - 1)
+        iterations_info = f"Planning Budget: Iteration {current_iteration + 1} of {max_iterations} (with {iterations_remaining} remaining)"
+
+        # Format the planning prompt
+        prompt = FULL_PLAN_PROMPT_TEMPLATE.format(
+            objective=objective,
+            plan_result=format_plan_result(plan_result),
+            plan_status=plan_status,
+            iterations_info=iterations_info,
+            agents=agents,
+        )
+
+        # Get structured response from LLM
+        try:
+            plan_msg = PromptMessageMultipart(
+                role="user", content=[TextContent(type="text", text=prompt)]
+            )
+            return await self._llm.structured([plan_msg], Plan, request_params)
+        except Exception as e:
+            self.logger.error(f"Failed to parse plan: {str(e)}")
+            return None
+
+    async def _get_next_step(
+        self, objective: str, plan_result: PlanResult, request_params: RequestParams
+    ) -> Optional[NextStep]:
+        """
+        Generate just the next step for iterative planning.
+
+        Args:
+            objective: The objective to achieve
+            plan_result: Current plan execution state
+            request_params: Request parameters
+
+        Returns:
+            Next step to execute, or None if parsing fails
+        """
+        # Format agent information
+        agents = "\n".join(
+            [self._format_agent_info(agent_name) for agent_name in self.agents.keys()]
+        )
+
+        # Determine plan status
+        if plan_result.is_complete:
+            plan_status = "Plan Status: Complete"
+        elif plan_result.step_results:
+            plan_status = "Plan Status: In Progress"
+        else:
+            plan_status = "Plan Status: Not Started"
+
+        # Calculate iteration information
+        max_iterations = request_params.max_iterations
+        current_iteration = len(plan_result.step_results)
+        iterations_remaining = max_iterations - current_iteration
+        iterations_info = (
+            f"Planning Budget: {iterations_remaining} of {max_iterations} iterations remaining"
+        )
+
+        # Format the planning prompt
+        prompt = ITERATIVE_PLAN_PROMPT_TEMPLATE.format(
+            objective=objective,
+            plan_result=format_plan_result(plan_result),
+            plan_status=plan_status,
+            iterations_info=iterations_info,
+            agents=agents,
+        )
+
+        # Get structured response from LLM
+        try:
+            plan_msg = PromptMessageMultipart(
+                role="user", content=[TextContent(type="text", text=prompt)]
+            )
+            return await self._llm.structured([plan_msg], NextStep, request_params)
+        except Exception as e:
+            self.logger.error(f"Failed to parse next step: {str(e)}")
+            return None
+
+    def _validate_agent_names(self, plan: Plan) -> None:
+        """
+        Validate all agent names in a plan before execution.
+
+        Args:
+            plan: The plan to validate
+        """
+        if plan is None:
+            self.logger.error("Cannot validate agent names: plan is None")
+            return
+
+        invalid_agents = []
+
+        for step in plan.steps:
+            for task in step.tasks:
+                if task.agent not in self.agents:
+                    invalid_agents.append(task.agent)
+
+        if invalid_agents:
+            available_agents = ", ".join(self.agents.keys())
+            invalid_list = ", ".join(invalid_agents)
+            self.logger.error(
+                f"Plan contains invalid agent names: {invalid_list}. Available agents: {available_agents}"
+            )
+
+    def _format_agent_info(self, agent_name: str) -> str:
+        """
+        Format agent information for display in prompts.
+
+        Args:
+            agent_name: Name of the agent to format
+
+        Returns:
+            Formatted agent information string
+        """
+        agent = self.agents.get(agent_name)
+        if not agent:
+            self.logger.error(f"Agent '{agent_name}' not found in orchestrator agents")
+            return ""
+
+        # Get agent instruction or default description
+        instruction = (
+            agent.instruction if hasattr(agent, "instruction") else f"Agent '{agent_name}'"
+        )
+
+        # Format with XML tags
+        return f'<fastagent:agent name="{agent_name}">{instruction}</fastagent:agent>'
+
+    async def _planner_generate_str(self, message: str, request_params: RequestParams) -> str:
+        """
+        Generate string response from the orchestrator's own LLM.
+
+        Args:
+            message: Message to send to the LLM
+            request_params: Request parameters
+
+        Returns:
+            String response from the LLM
+        """
+        # Create prompt message
+        prompt = PromptMessageMultipart(
+            role="user", content=[TextContent(type="text", text=message)]
+        )
+
+        # Get response from LLM
+        response = await self._llm.generate([prompt], request_params)
+        return response.all_text()
+
+    def _merge_request_params(self, request_params: Optional[RequestParams]) -> RequestParams:
+        """
+        Merge provided request parameters with defaults.
+
+        Args:
+            request_params: Optional request parameters to merge
+
+        Returns:
+            Merged request parameters
+        """
+        # Create orchestrator-specific defaults
+        defaults = RequestParams(
+            use_history=False,  # Orchestrator doesn't use history
+            max_iterations=5,  # Default to 5 iterations
+            maxTokens=8192,  # Higher limit for planning
+            parallel_tool_calls=True,
+        )
+
+        # If base params provided, merge with defaults
+        if request_params:
+            # Create copy of defaults
+            params = defaults.model_copy()
+            # Update with provided params
+            if isinstance(request_params, dict):
+                params = params.model_copy(update=request_params)
+            else:
+                params = params.model_copy(update=request_params.model_dump())
+
+            # Force specific settings
+            params.use_history = False
+            return params
+
+        return defaults

mcp_agent/{workflows/orchestrator → agents/workflow}/orchestrator_models.py

@@ -2,7 +2,7 @@ from typing import List
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from mcp_agent.workflows.orchestrator.orchestrator_prompts import (
+from mcp_agent.agents.workflow.orchestrator_prompts import (
     PLAN_RESULT_TEMPLATE,
     STEP_RESULT_TEMPLATE,
     TASK_RESULT_TEMPLATE,
@@ -50,21 +50,15 @@ class Plan(BaseModel):
         description="List of steps to execute sequentially",
         default_factory=list,
     )
-    is_complete: bool = Field(
-        description="Whether the overall plan objective is complete"
-    )
+    is_complete: bool = Field(description="Whether the overall plan objective is complete")
 
 
 class TaskWithResult(Task):
     """An individual task with its result"""
 
-    result: str = Field(
-        description="Result of executing the task", default="Task completed"
-    )
+    result: str = Field(description="Result of executing the task", default="Task completed")
 
-    agent: str = Field(
-        description="Name of the agent that executed this task", default=""
-    )
+    agent: str = Field(description="Name of the agent that executed this task", default="")
 
     model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
 
@@ -76,11 +70,9 @@ class StepResult(BaseModel):
     task_results: List[TaskWithResult] = Field(
         description="Results of executing each task", default_factory=list
     )
-    result: str = Field(
-        description="Result of executing the step", default="Step completed"
-    )
+    result: str = Field(description="Result of executing the step", default="Step completed")
 
-    def add_task_result(self, task_result: TaskWithResult):
+    def add_task_result(self, task_result: TaskWithResult) -> None:
         """Add a task result to this step"""
         if not isinstance(self.task_results, list):
             self.task_results = []
@@ -108,7 +100,7 @@ class PlanResult(BaseModel):
     result: str | None = None
     """Result of executing the plan"""
 
-    def add_step_result(self, step_result: StepResult):
+    def add_step_result(self, step_result: StepResult) -> None:
         """Add a step result to this plan"""
         if not isinstance(self.step_results, list):
             self.step_results = []
@@ -118,9 +110,7 @@ class PlanResult(BaseModel):
 class NextStep(Step):
     """Single next step in iterative planning"""
 
-    is_complete: bool = Field(
-        description="Whether the overall plan objective is complete"
-    )
+    is_complete: bool = Field(description="Whether the overall plan objective is complete")
 
 
 def format_task_result_text(task_result: TaskWithResult) -> str:
@@ -162,7 +152,7 @@ def format_plan_result_text(plan_result: PlanResult) -> str:
 
 def format_task_result_xml(task_result: TaskWithResult) -> str:
     """Format a task result with XML tags for better semantic understanding"""
-    from mcp_agent.workflows.llm.prompt_utils import format_fastagent_tag
+    from mcp_agent.llm.prompt_utils import format_fastagent_tag
 
     return format_fastagent_tag(
         "task-result",
@@ -178,7 +168,7 @@ def format_task_result_xml(task_result: TaskWithResult) -> str:
 
 def format_step_result_xml(step_result: StepResult) -> str:
     """Format a step result with XML tags for better semantic understanding"""
-    from mcp_agent.workflows.llm.prompt_utils import format_fastagent_tag
+    from mcp_agent.llm.prompt_utils import format_fastagent_tag
 
     # Format each task result with XML
     task_results = []
@@ -200,7 +190,7 @@ def format_step_result_xml(step_result: StepResult) -> str:
 
 def format_plan_result(plan_result: PlanResult) -> str:
     """Format the full plan execution state with XML for better semantic understanding"""
-    from mcp_agent.workflows.llm.prompt_utils import format_fastagent_tag
+    from mcp_agent.llm.prompt_utils import format_fastagent_tag
 
     # Format objective
     objective_tag = format_fastagent_tag("objective", plan_result.objective)