fast-agent-mcp 0.0.15__py3-none-any.whl → 0.1.0__py3-none-any.whl

mcp_agent/workflows/orchestrator/orchestrator.py

@@ -21,7 +21,7 @@ from mcp_agent.workflows.llm.augmented_llm import (
 )
 from mcp_agent.workflows.orchestrator.orchestrator_models import (
     format_plan_result,
-    format_step_result,
+    format_step_result_text,
     NextStep,
     Plan,
     PlanResult,
@@ -129,6 +129,27 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         self.logger = logger
         self.name = name
 
+        # Store agents by name - COMPLETE REWRITE OF AGENT STORAGE
+        self.agents = {}
+        for agent in available_agents:
+            # Fix: Remove all special handling of agent names and store them exactly as they are
+            agent_name = agent.name
+
+            # Verify if the name is actually "None" (string) or None (NoneType)
+            if agent_name == "None":
+                # Try to get a better name from config if available
+                if hasattr(agent, "config") and agent.config and agent.config.name:
+                    agent_name = agent.config.name
+            elif agent_name is None:
+                # Try to get a better name from config if available
+                if hasattr(agent, "config") and agent.config and agent.config.name:
+                    agent_name = agent.config.name
+                else:
+                    agent_name = f"unnamed_agent_{len(self.agents)}"
+
+            self.logger.info(f"Adding agent '{agent_name}' to orchestrator")
+            self.agents[agent_name] = agent
+
     async def generate(
         self,
         message: str | MessageParamT | List[MessageParamT],
@@ -163,15 +184,29 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         request_params: RequestParams | None = None,
     ) -> ModelT:
         """Request a structured LLM generation and return the result as a Pydantic model."""
+        import json
+        from pydantic import ValidationError
+
         params = self.get_request_params(request_params)
         result_str = await self.generate_str(message=message, request_params=params)
 
-        # Use AugmentedLLM's structured output handling
-        return await super().generate_structured(
-            message=result_str,
-            response_model=response_model,
-            request_params=params,
-        )
+        try:
+            # Directly parse JSON and create model instance
+            parsed_data = json.loads(result_str)
+            return response_model(**parsed_data)
+        except (json.JSONDecodeError, ValidationError) as e:
+            # Log the error and fall back to the original method if direct parsing fails
+            self.logger.error(
+                f"Direct JSON parsing failed: {str(e)}. Falling back to standard method."
+            )
+            self.logger.debug(f"Failed JSON content: {result_str}")
+
+            # Use AugmentedLLM's structured output handling as fallback
+            return await super().generate_structured(
+                message=result_str,
+                response_model=response_model,
+                request_params=params,
+            )
 
     async def execute(
         self, objective: str, request_params: RequestParams | None = None
@@ -205,11 +240,15 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
                 )
                 logger.debug(f"Iteration {iterations}: Iterative plan:", data=next_step)
                 plan = Plan(steps=[next_step], is_complete=next_step.is_complete)
+                # Validate agent names in the plan early
+                self._validate_agent_names(plan)
             elif self.plan_type == "full":
                 plan = await self._get_full_plan(
                     objective=objective, plan_result=plan_result, request_params=params
                 )
                 logger.debug(f"Iteration {iterations}: Full Plan:", data=plan)
+                # Validate agent names in the plan early
+                self._validate_agent_names(plan)
             else:
                 raise ValueError(f"Invalid plan type {self.plan_type}")
 
@@ -221,6 +260,7 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
                     plan_result.is_complete = True
 
                     # Synthesize final result into a single message
+                    # Use the structured XML format for better context
                     synthesis_prompt = SYNTHESIZE_PLAN_PROMPT_TEMPLATE.format(
                         plan_result=format_plan_result(plan_result)
                     )
@@ -265,6 +305,7 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         """Execute a step's subtasks in parallel and synthesize results"""
 
         step_result = StepResult(step=step, task_results=[])
+        # Use structured XML format for context to help agents better understand the context
         context = format_plan_result(previous_result)
 
         # Execute tasks
@@ -272,16 +313,18 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         error_tasks = []
 
         for task in step.tasks:
+            # Make sure we're using a valid agent name
             agent = self.agents.get(task.agent)
             if not agent:
-                # Instead of failing the entire step, track this as an error task
+                # Log a more prominent error - this is a serious problem that shouldn't happen
+                # with the improved prompt
                 self.logger.error(
-                    f"No agent found matching '{task.agent}'. Available agents: {list(self.agents.keys())}"
+                    f"AGENT VALIDATION ERROR: 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())}",
+                        f"Error: Agent '{task.agent}' not found. This indicates a problem with the plan generation. Available agents: {', '.join(self.agents.keys())}",
                     )
                 )
                 continue
@@ -292,8 +335,14 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
                 context=context,
             )
 
-            # All agents should now be LLM-capable
-            futures.append(agent._llm.generate_str(message=task_description))
+            # Handle both Agent objects and AugmentedLLM objects
+            from mcp_agent.workflows.llm.augmented_llm import AugmentedLLM
+
+            if isinstance(agent, AugmentedLLM):
+                futures.append(agent.generate_str(message=task_description))
+            else:
+                # Traditional Agent objects with _llm property
+                futures.append(agent._llm.generate_str(message=task_description))
 
         # Wait for all tasks (only if we have valid futures)
         results = await self.executor.execute(*futures) if futures else []
@@ -307,18 +356,29 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
 
             if task_index < len(results):
                 result = results[task_index]
-                step_result.add_task_result(
-                    TaskWithResult(**task.model_dump(), result=str(result))
+                # Create a TaskWithResult that includes the agent name for attribution
+                task_model = task.model_dump()
+                task_result = TaskWithResult(
+                    description=task_model["description"],
+                    agent=task_model["agent"],  # Track which agent produced this result
+                    result=str(result),
                 )
+                step_result.add_task_result(task_result)
                 task_index += 1
 
         # Add error task results
         for task, error_message in error_tasks:
+            task_model = task.model_dump()
             step_result.add_task_result(
-                TaskWithResult(**task.model_dump(), result=error_message)
+                TaskWithResult(
+                    description=task_model["description"],
+                    agent=task_model["agent"],
+                    result=f"ERROR: {error_message}",
+                )
             )
 
-        step_result.result = format_step_result(step_result)
+        # Use text formatting for display in logs
+        step_result.result = format_step_result_text(step_result)
         return step_result
 
     async def _get_full_plan(
@@ -328,30 +388,92 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         request_params: RequestParams | None = None,
     ) -> Plan:
         """Generate full plan considering previous results"""
+        import json
+        from pydantic import ValidationError
+        from mcp_agent.workflows.orchestrator.orchestrator_models import (
+            Plan,
+            Step,
+            AgentTask,
+        )
+
         params = self.get_request_params(request_params)
         params = params.model_copy(update={"use_history": False})
 
-        agents = "\n".join(
-            [
-                f"{idx}. {self._format_agent_info(agent)}"
-                for idx, agent in enumerate(self.agents, 1)
-            ]
-        )
+        # Debug: Print agent names before formatting
+        print("\n------ AGENT NAMES BEFORE FORMATTING ------")
+        for agent_name in self.agents.keys():
+            print(f"Agent name: '{agent_name}'")
+        print("------------------------------------------\n")
+
+        # Format agents without numeric prefixes for cleaner XML
+        agent_formats = []
+        for agent_name in self.agents.keys():
+            formatted = self._format_agent_info(agent_name)
+            agent_formats.append(formatted)
+            print(f"Formatted agent '{agent_name}':\n{formatted[:200]}...\n")
+
+        agents = "\n".join(agent_formats)
+
+        # Create clear plan status indicator for the template
+        plan_status = "Plan Status: Not Started"
+        if hasattr(plan_result, "is_complete"):
+            plan_status = (
+                "Plan Status: Complete"
+                if plan_result.is_complete
+                else "Plan Status: In Progress"
+            )
 
         prompt = FULL_PLAN_PROMPT_TEMPLATE.format(
             objective=objective,
             plan_result=format_plan_result(plan_result),
+            plan_status=plan_status,
             agents=agents,
         )
 
-        # Use planner directly - no verb manipulation needed
-        plan = await self.planner.generate_structured(
+        # Get raw JSON response from LLM
+        result_str = await self.planner.generate_str(
             message=prompt,
-            response_model=Plan,
             request_params=params,
         )
 
-        return plan
+        try:
+            # Parse JSON directly
+            data = json.loads(result_str)
+
+            # Create models manually to ensure agent names are preserved exactly as returned
+            steps = []
+            for step_data in data.get("steps", []):
+                tasks = []
+                for task_data in step_data.get("tasks", []):
+                    # Create AgentTask directly from dict, preserving exact agent string
+                    task = AgentTask(
+                        description=task_data.get("description", ""),
+                        agent=task_data.get("agent", ""),  # Preserve exact agent name
+                    )
+                    tasks.append(task)
+
+                # Create Step with the exact task objects we created
+                step = Step(description=step_data.get("description", ""), tasks=tasks)
+                steps.append(step)
+
+            # Create final Plan
+            plan = Plan(steps=steps, is_complete=data.get("is_complete", False))
+
+            return plan
+
+        except (json.JSONDecodeError, ValidationError, KeyError) as e:
+            # Log detailed error and fall back to the original method as last resort
+            self.logger.error(f"Error parsing plan JSON: {str(e)}")
+            self.logger.debug(f"Failed JSON content: {result_str}")
+
+            # Use the normal structured parsing as fallback
+            plan = await self.planner.generate_structured(
+                message=result_str,
+                response_model=Plan,
+                request_params=params,
+            )
+
+            return plan
 
     async def _get_next_step(
         self,
@@ -360,54 +482,134 @@ class Orchestrator(AugmentedLLM[MessageParamT, MessageT]):
         request_params: RequestParams | None = None,
     ) -> NextStep:
         """Generate just the next needed step"""
+        import json
+        from pydantic import ValidationError
+        from mcp_agent.workflows.orchestrator.orchestrator_models import (
+            NextStep,
+            AgentTask,
+        )
+
         params = self.get_request_params(request_params)
         params = params.model_copy(update={"use_history": False})
 
+        # Format agents without numeric prefixes for cleaner XML
+        # FIX: Iterate over agent names instead of agent objects
         agents = "\n".join(
-            [
-                f"{idx}. {self._format_agent_info(agent)}"
-                for idx, agent in enumerate(self.agents, 1)
-            ]
+            [self._format_agent_info(agent_name) for agent_name in self.agents.keys()]
         )
 
+        # Create clear plan status indicator for the template
+        plan_status = "Plan Status: Not Started"
+        if hasattr(plan_result, "is_complete"):
+            plan_status = (
+                "Plan Status: Complete"
+                if plan_result.is_complete
+                else "Plan Status: In Progress"
+            )
+
         prompt = ITERATIVE_PLAN_PROMPT_TEMPLATE.format(
             objective=objective,
             plan_result=format_plan_result(plan_result),
+            plan_status=plan_status,
             agents=agents,
         )
 
-        # Use planner directly - no verb manipulation needed
-        next_step = await self.planner.generate_structured(
+        # Get raw JSON response from LLM
+        result_str = await self.planner.generate_str(
             message=prompt,
-            response_model=NextStep,
             request_params=params,
         )
-        return next_step
+
+        try:
+            # Parse JSON directly
+            data = json.loads(result_str)
+
+            # Create task objects manually to preserve exact agent names
+            tasks = []
+            for task_data in data.get("tasks", []):
+                # Preserve the exact agent name as specified in the JSON
+                task = AgentTask(
+                    description=task_data.get("description", ""),
+                    agent=task_data.get("agent", ""),
+                )
+                tasks.append(task)
+
+            # Create step with manually constructed tasks
+            next_step = NextStep(
+                description=data.get("description", ""),
+                tasks=tasks,
+                is_complete=data.get("is_complete", False),
+            )
+
+            return next_step
+
+        except (json.JSONDecodeError, ValidationError, KeyError) as e:
+            # Log detailed error and fall back to the original method
+            self.logger.error(f"Error parsing next step JSON: {str(e)}")
+            self.logger.debug(f"Failed JSON content: {result_str}")
+
+            # Use the normal structured parsing as fallback
+            next_step = await self.planner.generate_structured(
+                message=result_str,
+                response_model=NextStep,
+                request_params=params,
+            )
+
+            return next_step
 
     def _format_server_info(self, server_name: str) -> str:
-        """Format server information for display to planners"""
+        """Format server information for display to planners using XML tags"""
+        from mcp_agent.workflows.llm.prompt_utils import format_server_info
+
         server_config = self.server_registry.get_server_config(server_name)
-        server_str = f"Server Name: {server_name}"
-        if not server_config:
-            return server_str
 
-        description = server_config.description
-        if description:
-            server_str = f"{server_str}\nDescription: {description}"
+        # Get description or empty string if not available
+        description = ""
+        if server_config and server_config.description:
+            description = server_config.description
+
+        return format_server_info(server_name, description)
+
+    def _validate_agent_names(self, plan: Plan) -> None:
+        """
+        Validate all agent names in a plan before execution.
+        This helps catch invalid agent references early.
+        """
+        invalid_agents = []
+
+        for step in plan.steps:
+            for task in step.tasks:
+                if task.agent not in self.agents:
+                    invalid_agents.append(task.agent)
 
-        return server_str
+        if invalid_agents:
+            available_agents = ", ".join(self.agents.keys())
+            invalid_list = ", ".join(invalid_agents)
+            error_msg = f"Plan contains invalid agent names: {invalid_list}. Available agents: {available_agents}"
+            self.logger.error(error_msg)
+            # We don't raise an exception here as the execution will handle invalid agents
+            # by logging errors for individual tasks
 
     def _format_agent_info(self, agent_name: str) -> str:
-        """Format Agent information for display to planners"""
+        """Format Agent information for display to planners using XML tags"""
+        from mcp_agent.workflows.llm.prompt_utils import format_agent_info
+
         agent = self.agents.get(agent_name)
         if not agent:
+            self.logger.error(f"Agent '{agent_name}' not found in orchestrator agents")
             return ""
+        instruction = agent.instruction
 
-        servers = "\n".join(
-            [
-                f"- {self._format_server_info(server_name)}"
-                for server_name in agent.server_names
-            ]
-        )
+        # Get servers information
+        server_info = []
+        for server_name in agent.server_names:
+            server_config = self.server_registry.get_server_config(server_name)
+            description = ""
+            if server_config and server_config.description:
+                description = server_config.description
+
+            server_info.append({"name": server_name, "description": description})
 
-        return f"Agent Name: {agent.name}\nDescription: {agent.instruction}\nServers in Agent: {servers}"
+        return format_agent_info(
+            agent.name, instruction, server_info if server_info else None
+        )

mcp_agent/workflows/orchestrator/orchestrator_models.py

@@ -62,6 +62,10 @@ class TaskWithResult(Task):
         description="Result of executing the task", default="Task completed"
     )
 
+    agent: str = Field(
+        description="Name of the agent that executed this task", default=""
+    )
+
     model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
 
 
@@ -116,17 +120,17 @@ class NextStep(Step):
     )
 
 
-def format_task_result(task_result: TaskWithResult) -> str:
-    """Format a task result for display to planners"""
+def format_task_result_text(task_result: TaskWithResult) -> str:
+    """Format a task result as plain text for display"""
     return TASK_RESULT_TEMPLATE.format(
         task_description=task_result.description, task_result=task_result.result
     )
 
 
-def format_step_result(step_result: StepResult) -> str:
-    """Format a step result for display to planners"""
+def format_step_result_text(step_result: StepResult) -> str:
+    """Format a step result as plain text for display"""
     tasks_str = "\n".join(
-        f"  - {format_task_result(task)}" for task in step_result.task_results
+        f"  - {format_task_result_text(task)}" for task in step_result.task_results
     )
     return STEP_RESULT_TEMPLATE.format(
         step_description=step_result.step.description,
@@ -135,11 +139,11 @@ def format_step_result(step_result: StepResult) -> str:
     )
 
 
-def format_plan_result(plan_result: PlanResult) -> str:
-    """Format the full plan execution state for display to planners"""
+def format_plan_result_text(plan_result: PlanResult) -> str:
+    """Format the full plan execution state as plain text for display"""
     steps_str = (
         "\n\n".join(
-            f"{i + 1}:\n{format_step_result(step)}"
+            f"{i + 1}:\n{format_step_result_text(step)}"
             for i, step in enumerate(plan_result.step_results)
         )
         if plan_result.step_results
@@ -149,6 +153,74 @@ def format_plan_result(plan_result: PlanResult) -> str:
     return PLAN_RESULT_TEMPLATE.format(
         plan_objective=plan_result.objective,
         steps_str=steps_str,
-        plan_status="Complete" if plan_result.is_complete else "In Progress",
         plan_result=plan_result.result if plan_result.is_complete else "In Progress",
     )
+
+
+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
+
+    return format_fastagent_tag(
+        "task-result",
+        f"\n<fastagent:description>{task_result.description}</fastagent:description>\n"
+        f"<fastagent:result>{task_result.result}</fastagent:result>\n",
+        {
+            "description": task_result.description[:50] + "..."
+            if len(task_result.description) > 50
+            else task_result.description
+        },
+    )
+
+
+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
+
+    # Format each task result with XML
+    task_results = []
+    for task in step_result.task_results:
+        task_results.append(format_task_result_xml(task))
+
+    # Combine task results
+    task_results_str = "\n".join(task_results)
+
+    # Build step result with metadata and tasks
+    step_content = (
+        f"<fastagent:description>{step_result.step.description}</fastagent:description>\n"
+        f"<fastagent:summary>{step_result.result}</fastagent:summary>\n"
+        f"<fastagent:task-results>\n{task_results_str}\n</fastagent:task-results>\n"
+    )
+
+    return format_fastagent_tag("step-result", step_content)
+
+
+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
+
+    # Format objective
+    objective_tag = format_fastagent_tag("objective", plan_result.objective)
+
+    # Format step results
+    step_results = []
+    for step in plan_result.step_results:
+        step_results.append(format_step_result_xml(step))
+
+    # Build progress section
+    if step_results:
+        steps_content = "\n".join(step_results)
+        progress_content = (
+            f"{objective_tag}\n"
+            f"<fastagent:steps>\n{steps_content}\n</fastagent:steps>\n"
+            f"<fastagent:status>{plan_result.result if plan_result.is_complete else 'In Progress'}</fastagent:status>\n"
+        )
+    else:
+        # No steps executed yet
+        progress_content = (
+            f"{objective_tag}\n"
+            f"<fastagent:steps>No steps executed yet</fastagent:steps>\n"
+            f"<fastagent:status>Not Started</fastagent:status>\n"
+        )
+
+    return format_fastagent_tag("progress", progress_content)