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

mcp_agent/workflows/llm/augmented_llm_anthropic.py

@@ -117,10 +117,11 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         responses: List[Message] = []
         model = await self.select_model(params)
         chat_turn = (len(messages) + 1) // 2
-        self._log_chat_progress(chat_turn, model=model)
         self.show_user_message(str(message), model, chat_turn)
 
         for i in range(params.max_iterations):
+            chat_turn = (len(messages) + 1) // 2
+            self._log_chat_progress(chat_turn, model=model)
             arguments = {
                 "model": model,
                 "messages": messages,
@@ -208,10 +209,23 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
                 break
             elif response.stop_reason == "max_tokens":
                 # We have reached the max tokens limit
+
                 self.logger.debug(
                     f"Iteration {i}: Stopping because finish_reason is 'max_tokens'"
                 )
-                # TODO: saqadri - would be useful to return the reason for stopping to the caller
+                if params.maxTokens is not None:
+                    message_text = Text(
+                        f"the assistant has reached the maximum token limit ({params.maxTokens})",
+                        style="dim green italic",
+                    )
+                else:
+                    message_text = Text(
+                        "the assistant has reached the maximum token limit",
+                        style="dim green italic",
+                    )
+
+                await self.show_assistant_message(message_text)
+
                 break
             else:
                 message_text = ""

mcp_agent/workflows/llm/augmented_llm_openai.py

@@ -254,7 +254,7 @@ class OpenAIAugmentedLLM(
                         message_text,
                         message.tool_calls[
                             0
-                        ].function.name,  # TODO support multiple tool calls
+                        ].function.name,  # TODO support displaying multiple tool calls
                     )
                 else:
                     await self.show_assistant_message(
@@ -294,6 +294,18 @@ class OpenAIAugmentedLLM(
                 self.logger.debug(
                     f"Iteration {i}: Stopping because finish_reason is 'length'"
                 )
+                if request_params and request_params.maxTokens is not None:
+                    message_text = Text(
+                        f"the assistant has reached the maximum token limit ({request_params.maxTokens})",
+                        style="dim green italic",
+                    )
+                else:
+                    message_text = Text(
+                        "the assistant has reached the maximum token limit",
+                        style="dim green italic",
+                    )
+
+                await self.show_assistant_message(message_text)
                 # TODO: saqadri - would be useful to return the reason for stopping to the caller
                 break
             elif choice.finish_reason == "content_filter":

mcp_agent/workflows/llm/prompt_utils.py

@@ -0,0 +1,137 @@
+"""
+XML formatting utilities for consistent prompt engineering across components.
+"""
+
+from typing import Dict, List, Optional, Union
+
+
+def format_xml_tag(tag_name: str, content: Optional[str] = None, 
+                 attributes: Optional[Dict[str, str]] = None) -> str:
+    """
+    Format an XML tag with optional content and attributes.
+    Uses self-closing tag when content is None or empty.
+
+    Args:
+        tag_name: Name of the XML tag
+        content: Content to include inside the tag (None for self-closing)
+        attributes: Dictionary of attribute name-value pairs
+
+    Returns:
+        Formatted XML tag as string
+    """
+    # Format attributes if provided
+    attrs_str = ""
+    if attributes:
+        attrs_str = " " + " ".join(f'{k}="{v}"' for k, v in attributes.items())
+    
+    # Use self-closing tag if no content
+    if content is None or content == "":
+        return f"<{tag_name}{attrs_str} />"
+    
+    # Full tag with content
+    return f"<{tag_name}{attrs_str}>{content}</{tag_name}>"
+
+
+def format_fastagent_tag(tag_type: str, content: Optional[str] = None, 
+                        attributes: Optional[Dict[str, str]] = None) -> str:
+    """
+    Format a fastagent-namespaced XML tag with consistent formatting.
+    
+    Args:
+        tag_type: Type of fastagent tag (without namespace prefix)
+        content: Content to include inside the tag
+        attributes: Dictionary of attribute name-value pairs
+        
+    Returns:
+        Formatted fastagent XML tag as string
+    """
+    return format_xml_tag(f"fastagent:{tag_type}", content, attributes)
+
+
+def format_server_info(server_name: str, description: Optional[str] = None, 
+                      tools: Optional[List[Dict[str, str]]] = None) -> str:
+    """
+    Format server information consistently across router and orchestrator modules.
+    
+    Args:
+        server_name: Name of the server
+        description: Optional server description
+        tools: Optional list of tool dictionaries with 'name' and 'description' keys
+        
+    Returns:
+        Formatted server XML as string
+    """
+    # Use self-closing tag if no description or tools
+    if not description and not tools:
+        return format_fastagent_tag("server", None, {"name": server_name})
+    
+    # Start building components
+    components = []
+    
+    # Add description if present
+    if description:
+        desc_tag = format_fastagent_tag("description", description)
+        components.append(desc_tag)
+    
+    # Add tools section if tools exist
+    if tools and len(tools) > 0:
+        tool_tags = []
+        for tool in tools:
+            tool_name = tool.get("name", "")
+            tool_desc = tool.get("description", "")
+            tool_tag = format_fastagent_tag("tool", tool_desc, {"name": tool_name})
+            tool_tags.append(tool_tag)
+        
+        tools_content = "\n".join(tool_tags)
+        tools_tag = format_fastagent_tag("tools", f"\n{tools_content}\n")
+        components.append(tools_tag)
+    
+    # Combine all components
+    server_content = "\n".join(components)
+    return format_fastagent_tag("server", f"\n{server_content}\n", {"name": server_name})
+
+
+def format_agent_info(agent_name: str, description: Optional[str] = None, 
+                     servers: Optional[List[Dict[str, Union[str, List[Dict[str, str]]]]]] = None) -> str:
+    """
+    Format agent information consistently across router and orchestrator modules.
+    
+    Args:
+        agent_name: Name of the agent
+        description: Optional agent description/instruction
+        servers: Optional list of server dictionaries with 'name', 'description', and 'tools' keys
+        
+    Returns:
+        Formatted agent XML as string
+    """
+    # Start building components
+    components = []
+    
+    # Add description if present
+    if description:
+        desc_tag = format_fastagent_tag("description", description)
+        components.append(desc_tag)
+    
+    # If no description or servers, use self-closing tag
+    if not description and not servers:
+        return format_fastagent_tag("agent", None, {"name": agent_name})
+        
+    # If has servers, format them
+    if servers and len(servers) > 0:
+        server_tags = []
+        for server in servers:
+            server_name = server.get("name", "")
+            server_desc = server.get("description", "")
+            server_tools = server.get("tools", [])
+            server_tag = format_server_info(server_name, server_desc, server_tools)
+            server_tags.append(server_tag)
+        
+        # Only add servers section if we have servers
+        if server_tags:
+            servers_content = "\n".join(server_tags)
+            servers_tag = format_fastagent_tag("servers", f"\n{servers_content}\n")
+            components.append(servers_tag)
+    
+    # Combine all components
+    agent_content = "\n".join(components)
+    return format_fastagent_tag("agent", f"\n{agent_content}\n", {"name": agent_name})

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,
@@ -163,15 +163,27 @@ 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 +217,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 +237,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 +282,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 +290,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
@@ -307,18 +327,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 +359,80 @@ 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})
 
+        # Format agents without numeric prefixes for cleaner XML
         agents = "\n".join(
-            [
-                f"{idx}. {self._format_agent_info(agent)}"
-                for idx, agent in enumerate(self.agents, 1)
-            ]
+            [self._format_agent_info(agent) for agent in self.agents]
         )
 
+        # 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 +441,127 @@ 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
         agents = "\n".join(
-            [
-                f"{idx}. {self._format_agent_info(agent)}"
-                for idx, agent in enumerate(self.agents, 1)
-            ]
+            [self._format_agent_info(agent) for agent in self.agents]
         )
 
+        # 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}"
-
-        return server_str
-
+        
+        # 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)
+        
+        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:
             return ""
 
-        servers = "\n".join(
-            [
-                f"- {self._format_server_info(server_name)}"
-                for server_name in agent.server_names
-            ]
-        )
-
-        return f"Agent Name: {agent.name}\nDescription: {agent.instruction}\nServers in Agent: {servers}"
+        # Get agent instruction as string
+        instruction = agent.instruction
+        if callable(instruction):
+            instruction = instruction({})
+        
+        # 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 format_agent_info(agent.name, instruction, server_info if server_info else None)