fast-agent-mcp 0.0.7__py3-none-any.whl

mcp_agent/workflows/router/router_llm.py

@@ -0,0 +1,301 @@
+from typing import Callable, List, Literal, Optional, TYPE_CHECKING
+
+from pydantic import BaseModel
+
+from mcp_agent.agents.agent import Agent
+from mcp_agent.workflows.llm.augmented_llm import AugmentedLLM, RequestParams
+from mcp_agent.workflows.router.router_base import ResultT, Router, RouterResult
+from mcp_agent.logging.logger import get_logger
+
+if TYPE_CHECKING:
+    from mcp_agent.context import Context
+
+logger = get_logger(__name__)
+
+
+DEFAULT_ROUTING_INSTRUCTION = """
+You are a highly accurate request router that directs incoming requests to the most appropriate category.
+A category is a specialized destination, such as a Function, an MCP Server (a collection of tools/functions), or an Agent (a collection of servers).
+Below are the available routing categories, each with their capabilities and descriptions:
+
+{context}
+
+Your task is to analyze the following request and determine the most appropriate categories from the options above. Consider:
+- The specific capabilities and tools each destination offers
+- How well the request matches the category's description
+- Whether the request might benefit from multiple categories (up to {top_k})
+
+Request: {request}
+
+Respond in JSON format:
+{{
+    "categories": [
+        {{
+            "category": <category name>,
+            "confidence": <high, medium or low>,
+            "reasoning": <brief explanation>
+        }}
+    ]
+}}
+
+Only include categories that are truly relevant. You may return fewer than {top_k} if appropriate.
+If none of the categories are relevant, return an empty list.
+"""
+
+ROUTING_SYSTEM_INSTRUCTION = """
+You are a highly accurate request router that directs incoming requests to the most appropriate category.
+A category is a specialized destination, such as a Function, an MCP Server (a collection of tools/functions), or an Agent (a collection of servers).
+You will be provided with a request and a list of categories to choose from.
+You can choose one or more categories, or choose none if no category is appropriate.
+"""
+
+
+class LLMRouterResult(RouterResult[ResultT]):
+    """A class that represents the result of an LLMRouter.route request"""
+
+    confidence: Literal["high", "medium", "low"]
+    """The confidence level of the routing decision."""
+
+    reasoning: str | None = None
+    """
+    A brief explanation of the routing decision.
+    This is optional and may only be provided if the router is an LLM
+    """
+
+
+class StructuredResponseCategory(BaseModel):
+    """A class that represents a single category returned by an LLM router"""
+
+    category: str
+    """The name of the category (i.e. MCP server, Agent or function) to route the input to."""
+
+    confidence: Literal["high", "medium", "low"]
+    """The confidence level of the routing decision."""
+
+    reasoning: str | None = None
+    """A brief explanation of the routing decision."""
+
+
+class StructuredResponse(BaseModel):
+    """A class that represents the structured response of an LLM router"""
+
+    categories: List[StructuredResponseCategory]
+    """A list of categories to route the input to."""
+
+
+class LLMRouter(Router):
+    """
+    A router that uses an LLM to route an input to a specific category.
+    """
+
+    def __init__(
+        self,
+        llm_factory: Callable[..., AugmentedLLM],
+        name: str = "LLM Router",
+        server_names: List[str] | None = None,
+        agents: List[Agent] | None = None,
+        functions: List[Callable] | None = None,
+        routing_instruction: str | None = None,
+        context: Optional["Context"] = None,
+        default_request_params: Optional[RequestParams] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            server_names=server_names,
+            agents=agents,
+            functions=functions,
+            routing_instruction=routing_instruction,
+            context=context,
+            **kwargs,
+        )
+
+        self.name = name
+        self.llm_factory = llm_factory
+        self.default_request_params = default_request_params or RequestParams()
+        self.llm = None  # Will be initialized in create()
+
+    @classmethod
+    async def create(
+        cls,
+        llm_factory: Callable[..., AugmentedLLM],
+        name: str = "LLM Router",
+        server_names: List[str] | None = None,
+        agents: List[Agent] | None = None,
+        functions: List[Callable] | None = None,
+        routing_instruction: str | None = None,
+        context: Optional["Context"] = None,
+        default_request_params: Optional[RequestParams] = None,
+    ) -> "LLMRouter":
+        """
+        Factory method to create and initialize a router.
+        Use this instead of constructor since we need async initialization.
+        """
+        instance = cls(
+            llm_factory=llm_factory,
+            name=name,
+            server_names=server_names,
+            agents=agents,
+            functions=functions,
+            routing_instruction=DEFAULT_ROUTING_INSTRUCTION,
+            context=context,
+            default_request_params=default_request_params,
+        )
+        await instance.initialize()
+        return instance
+
+    async def initialize(self):
+        """Initialize the router and create the LLM instance."""
+        if not self.initialized:
+            await super().initialize()
+            router_params = RequestParams(
+                systemPrompt=ROUTING_SYSTEM_INSTRUCTION,
+                use_history=False,  # Router should be stateless :)
+            )
+
+            # Merge with any provided default params
+            if self.default_request_params:
+                params_dict = router_params.model_dump()
+                params_dict.update(
+                    self.default_request_params.model_dump(exclude_unset=True)
+                )
+                router_params = RequestParams(**params_dict)
+            # Set up router-specific request params with routing instruction
+            router_params.use_history = False
+            self.llm = self.llm_factory(
+                agent=None,  # Router doesn't need an agent context
+                name="LLM Router",
+                default_request_params=router_params,
+            )
+            self.initialized = True
+
+    async def route(
+        self, request: str, top_k: int = 1
+    ) -> List[LLMRouterResult[str | Agent | Callable]]:
+        if not self.initialized:
+            await self.initialize()
+
+        return await self._route_with_llm(request, top_k)
+
+    async def route_to_server(
+        self, request: str, top_k: int = 1
+    ) -> List[LLMRouterResult[str]]:
+        if not self.initialized:
+            await self.initialize()
+
+        return await self._route_with_llm(
+            request,
+            top_k,
+            include_servers=True,
+            include_agents=False,
+            include_functions=False,
+        )
+
+    async def route_to_agent(
+        self, request: str, top_k: int = 1
+    ) -> List[LLMRouterResult[Agent]]:
+        if not self.initialized:
+            await self.initialize()
+
+        return await self._route_with_llm(
+            request,
+            top_k,
+            include_servers=False,
+            include_agents=True,
+            include_functions=False,
+        )
+
+    async def route_to_function(
+        self, request: str, top_k: int = 1
+    ) -> List[LLMRouterResult[Callable]]:
+        if not self.initialized:
+            await self.initialize()
+
+        return await self._route_with_llm(
+            request,
+            top_k,
+            include_servers=False,
+            include_agents=False,
+            include_functions=True,
+        )
+
+    async def _route_with_llm(
+        self,
+        request: str,
+        top_k: int = 1,
+        include_servers: bool = True,
+        include_agents: bool = True,
+        include_functions: bool = True,
+    ) -> List[LLMRouterResult]:
+        if not self.initialized:
+            await self.initialize()
+
+        routing_instruction = self.routing_instruction or DEFAULT_ROUTING_INSTRUCTION
+
+        # Generate the categories context
+        context = self._generate_context(
+            include_servers=include_servers,
+            include_agents=include_agents,
+            include_functions=include_functions,
+        )
+
+        # Format the prompt with all the necessary information
+        prompt = routing_instruction.format(
+            context=context, request=request, top_k=top_k
+        )
+
+        # Get routes from LLM
+        response = await self.llm.generate_structured(
+            message=prompt,
+            response_model=StructuredResponse,
+        )
+
+        # Construct the result
+        if not response or not response.categories:
+            return []
+
+        result: List[LLMRouterResult] = []
+        for r in response.categories:
+            router_category = self.categories.get(r.category)
+            if not router_category:
+                # Skip invalid categories
+                # TODO: log or raise an error
+                continue
+
+            result.append(
+                LLMRouterResult(
+                    result=router_category.category,
+                    confidence=r.confidence,
+                    reasoning=r.reasoning,
+                )
+            )
+
+        return result[:top_k]
+
+    def _generate_context(
+        self,
+        include_servers: bool = True,
+        include_agents: bool = True,
+        include_functions: bool = True,
+    ) -> str:
+        """Generate a formatted context list of categories."""
+
+        context_list = []
+        idx = 1
+
+        # Format all categories
+        if include_servers:
+            for category in self.server_categories.values():
+                context_list.append(self.format_category(category, idx))
+                idx += 1
+
+        if include_agents:
+            for category in self.agent_categories.values():
+                context_list.append(self.format_category(category, idx))
+                idx += 1
+
+        if include_functions:
+            for category in self.function_categories.values():
+                context_list.append(self.format_category(category, idx))
+                idx += 1
+
+        return "\n\n".join(context_list)

mcp_agent/workflows/swarm/swarm.py

@@ -0,0 +1,320 @@
+from typing import Callable, Dict, Generic, List, Optional, TYPE_CHECKING
+from collections import defaultdict
+
+from pydantic import AnyUrl, BaseModel, ConfigDict
+from mcp.types import (
+    CallToolRequest,
+    EmbeddedResource,
+    CallToolResult,
+    TextContent,
+    TextResourceContents,
+    Tool,
+)
+
+from mcp_agent.agents.agent import Agent
+from mcp_agent.human_input.types import HumanInputCallback
+from mcp_agent.workflows.llm.augmented_llm import (
+    AugmentedLLM,
+    MessageParamT,
+    MessageT,
+)
+from mcp_agent.logging.logger import get_logger
+
+if TYPE_CHECKING:
+    from mcp_agent.context import Context
+
+logger = get_logger(__name__)
+
+
+class AgentResource(EmbeddedResource):
+    """
+    A resource that returns an agent. Meant for use with tool calls that want to return an Agent for further processing.
+    """
+
+    agent: Optional["Agent"] = None
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+
+class AgentFunctionResultResource(EmbeddedResource):
+    """
+    A resource that returns an AgentFunctionResult.
+    Meant for use with tool calls that return an AgentFunctionResult for further processing.
+    """
+
+    result: "AgentFunctionResult"
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+
+def create_agent_resource(agent: "Agent") -> AgentResource:
+    return AgentResource(
+        type="resource",
+        agent=agent,
+        resource=TextResourceContents(
+            text=f"You are now Agent '{agent.name}'. Please review the messages and continue execution",
+            uri=AnyUrl("http://fake.url"),  # Required property but not needed
+        ),
+    )
+
+
+def create_agent_function_result_resource(
+    result: "AgentFunctionResult",
+) -> AgentFunctionResultResource:
+    return AgentFunctionResultResource(
+        type="resource",
+        result=result,
+        resource=TextResourceContents(
+            text=result.value or result.agent.name or "AgentFunctionResult",
+            uri=AnyUrl("http://fake.url"),  # Required property but not needed
+        ),
+    )
+
+
+class SwarmAgent(Agent):
+    """
+    A SwarmAgent is an Agent that can spawn other agents and interactively resolve a task.
+    Based on OpenAI Swarm: https://github.com/openai/swarm.
+
+    SwarmAgents have access to tools available on the servers they are connected to, but additionally
+    have a list of (possibly local) functions that can be called as tools.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        instruction: str | Callable[[Dict], str] = "You are a helpful agent.",
+        server_names: list[str] = None,
+        functions: List["AgentFunctionCallable"] = None,
+        parallel_tool_calls: bool = True,
+        human_input_callback: HumanInputCallback = None,
+        context: Optional["Context"] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            instruction=instruction,
+            server_names=server_names,
+            functions=functions,
+            # TODO: saqadri - figure out if Swarm can maintain connection persistence
+            # It's difficult because we don't know when the agent will be done with its task
+            connection_persistence=False,
+            human_input_callback=human_input_callback,
+            context=context,
+            **kwargs,
+        )
+        self.parallel_tool_calls = parallel_tool_calls
+
+    async def call_tool(
+        self, name: str, arguments: dict | None = None
+    ) -> CallToolResult:
+        if not self.initialized:
+            await self.initialize()
+
+        if name in self._function_tool_map:
+            tool = self._function_tool_map[name]
+            result = await tool.run(arguments)
+
+            logger.debug(f"Function tool {name} result:", data=result)
+
+            if isinstance(result, Agent) or isinstance(result, SwarmAgent):
+                resource = create_agent_resource(result)
+                return CallToolResult(content=[resource])
+            elif isinstance(result, AgentFunctionResult):
+                resource = create_agent_function_result_resource(result)
+                return CallToolResult(content=[resource])
+            elif isinstance(result, str):
+                # TODO: saqadri - this is likely meant for returning context variables
+                return CallToolResult(content=[TextContent(type="text", text=result)])
+            elif isinstance(result, dict):
+                return CallToolResult(
+                    content=[TextContent(type="text", text=str(result))]
+                )
+            else:
+                logger.warning(f"Unknown result type: {result}, returning as text.")
+                return CallToolResult(
+                    content=[TextContent(type="text", text=str(result))]
+                )
+
+        return await super().call_tool(name, arguments)
+
+
+class AgentFunctionResult(BaseModel):
+    """
+    Encapsulates the possible return values for a Swarm agent function.
+
+    Attributes:
+        value (str): The result value as a string.
+        agent (Agent): The agent instance, if applicable.
+        context_variables (dict): A dictionary of context variables.
+    """
+
+    value: str = ""
+    agent: Agent | None = None
+    context_variables: dict = {}
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+
+AgentFunctionReturnType = str | Agent | dict | AgentFunctionResult
+"""A type alias for the return type of a Swarm agent function."""
+
+AgentFunctionCallable = Callable[[], AgentFunctionReturnType]
+
+
+async def create_transfer_to_agent_tool(
+    agent: "Agent", agent_function: Callable[[], None]
+) -> Tool:
+    return Tool(
+        name="transfer_to_agent",
+        description="Transfer control to the agent",
+        agent_resource=create_agent_resource(agent),
+        agent_function=agent_function,
+    )
+
+
+async def create_agent_function_tool(agent_function: "AgentFunctionCallable") -> Tool:
+    return Tool(
+        name="agent_function",
+        description="Agent function",
+        agent_resource=None,
+        agent_function=agent_function,
+    )
+
+
+class Swarm(AugmentedLLM[MessageParamT, MessageT], Generic[MessageParamT, MessageT]):
+    """
+    Handles orchestrating agents that can use tools via MCP servers.
+
+    MCP version of the OpenAI Swarm class (https://github.com/openai/swarm.)
+    """
+
+    # TODO: saqadri - streaming isn't supported yet because the underlying AugmentedLLM classes don't support it
+    def __init__(self, agent: SwarmAgent, context_variables: Dict[str, str] = None):
+        """
+        Initialize the LLM planner with an agent, which will be used as the
+        starting point for the workflow.
+        """
+        super().__init__(agent=agent)
+        self.context_variables = defaultdict(str, context_variables or {})
+        self.instruction = (
+            agent.instruction(self.context_variables)
+            if isinstance(agent.instruction, Callable)
+            else agent.instruction
+        )
+        logger.debug(
+            f"Swarm initialized with agent {agent.name}",
+            data={
+                "context_variables": self.context_variables,
+                "instruction": self.instruction,
+            },
+        )
+
+    async def get_tool(self, tool_name: str) -> Tool | None:
+        """Get the schema for a tool by name."""
+        result = await self.aggregator.list_tools()
+        for tool in result.tools:
+            if tool.name == tool_name:
+                return tool
+
+        return None
+
+    async def pre_tool_call(
+        self, tool_call_id: str | None, request: CallToolRequest
+    ) -> CallToolRequest | bool:
+        if not self.aggregator:
+            # If there are no agents, we can't do anything, so we should bail
+            return False
+
+        tool = await self.get_tool(request.params.name)
+        if not tool:
+            logger.warning(
+                f"Warning: Tool '{request.params.name}' not found in agent '{self.aggregator.name}' tools. Proceeding with original request params."
+            )
+            return request
+
+        # If the tool has a "context_variables" parameter, we set it to our context variables state
+        if "context_variables" in tool.inputSchema:
+            logger.debug(
+                f"Setting context variables on tool_call '{request.params.name}'",
+                data=self.context_variables,
+            )
+            request.params.arguments["context_variables"] = self.context_variables
+
+        return request
+
+    async def post_tool_call(
+        self, tool_call_id: str | None, request: CallToolRequest, result: CallToolResult
+    ) -> CallToolResult:
+        contents = []
+        for content in result.content:
+            if isinstance(content, AgentResource):
+                # Set the new agent as the current agent
+                await self.set_agent(content.agent)
+                contents.append(TextContent(type="text", text=content.resource.text))
+            elif isinstance(content, AgentFunctionResult):
+                logger.info(
+                    "Updating context variables with new context variables from agent function result",
+                    data=content.context_variables,
+                )
+                self.context_variables.update(content.context_variables)
+                if content.agent:
+                    # Set the new agent as the current agent
+                    self.set_agent(content.agent)
+
+                contents.append(TextContent(type="text", text=content.resource.text))
+            else:
+                contents.append(content)
+
+        result.content = contents
+        return result
+
+    async def set_agent(
+        self,
+        agent: SwarmAgent,
+    ):
+        logger.info(
+            f"Switching from agent '{self.aggregator.name}' -> agent '{agent.name if agent else 'NULL'}'"
+        )
+        if self.aggregator:
+            # Close the current agent
+            await self.aggregator.shutdown()
+
+        # Initialize the new agent (if it's not None)
+        self.aggregator = agent
+
+        if not self.aggregator or isinstance(self.aggregator, DoneAgent):
+            self.instruction = None
+            return
+
+        await self.aggregator.initialize()
+        self.instruction = (
+            agent.instruction(self.context_variables)
+            if callable(agent.instruction)
+            else agent.instruction
+        )
+
+    def should_continue(self) -> bool:
+        """
+        Returns True if the workflow should continue, False otherwise.
+        """
+        if not self.aggregator or isinstance(self.aggregator, DoneAgent):
+            return False
+
+        return True
+
+
+class DoneAgent(SwarmAgent):
+    """
+    A special agent that represents the end of a Swarm workflow.
+    """
+
+    def __init__(self):
+        super().__init__(name="__done__", instruction="Swarm Workflow is complete.")
+
+    async def call_tool(
+        self, _name: str, _arguments: dict | None = None
+    ) -> CallToolResult:
+        return CallToolResult(
+            content=[TextContent(type="text", text="Workflow is complete.")]
+        )

mcp_agent/workflows/swarm/swarm_anthropic.py

@@ -0,0 +1,42 @@
+from mcp_agent.workflows.swarm.swarm import Swarm
+from mcp_agent.workflows.llm.augmented_llm import RequestParams
+from mcp_agent.workflows.llm.augmented_llm_anthropic import AnthropicAugmentedLLM
+from mcp_agent.logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class AnthropicSwarm(Swarm, AnthropicAugmentedLLM):
+    """
+    MCP version of the OpenAI Swarm class (https://github.com/openai/swarm.),
+    using Anthropic's API as the LLM.
+    """
+
+    async def generate(self, message, request_params: RequestParams | None = None):
+        params = self.get_request_params(
+            request_params,
+            default=RequestParams(
+                model="claude-3-5-sonnet-20241022",
+                maxTokens=8192,
+                parallel_tool_calls=False,
+            ),
+        )
+        iterations = 0
+        response = None
+        agent_name = str(self.aggregator.name) if self.aggregator else None
+
+        while iterations < params.max_iterations and self.should_continue():
+            response = await super().generate(
+                message=message
+                if iterations == 0
+                else "Please resolve my original request. If it has already been resolved then end turn",
+                request_params=params.model_copy(
+                    update={"max_iterations": 1}
+                ),  # TODO: saqadri - validate
+            )
+            logger.debug(f"Agent: {agent_name}, response:", data=response)
+            agent_name = self.aggregator.name if self.aggregator else None
+            iterations += 1
+
+        # Return final response back
+        return response