rasa-pro 3.14.0.dev5__py3-none-any.whl → 3.14.0.dev6__py3-none-any.whl

rasa/agents/protocol/mcp/mcp_base_agent.py

@@ -0,0 +1,684 @@
+import json
+from abc import abstractmethod
+from inspect import isawaitable
+from typing import Any, Dict, List, Optional, Tuple
+
+import structlog
+from jinja2 import Template
+from mcp import ListToolsResult
+
+from rasa.agents.constants import (
+    AGENT_DEFAULT_MAX_RETRIES,
+    AGENT_DEFAULT_TIMEOUT_SECONDS,
+    AGENT_METADATA_TOOL_RESULTS_KEY,
+    KEY_ARGUMENTS,
+    KEY_CONTENT,
+    KEY_FUNCTION,
+    KEY_ID,
+    KEY_NAME,
+    KEY_ROLE,
+    KEY_TOOL_CALL_ID,
+    KEY_TOOL_CALLS,
+    KEY_TYPE,
+)
+from rasa.agents.core.agent_protocol import AgentProtocol
+from rasa.agents.core.types import AgentStatus, ProtocolType
+from rasa.agents.schemas import (
+    AgentInput,
+    AgentOutput,
+    AgentToolResult,
+    AgentToolSchema,
+    CustomToolSchema,
+)
+from rasa.core.available_agents import AgentConfig, AgentMCPServerConfig, ProtocolConfig
+from rasa.shared.agents.utils import make_agent_identifier
+from rasa.shared.constants import (
+    MAX_COMPLETION_TOKENS_CONFIG_KEY,
+    MODEL_CONFIG_KEY,
+    OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
+    ROLE_ASSISTANT,
+    ROLE_SYSTEM,
+    ROLE_TOOL,
+    ROLE_USER,
+    TEMPERATURE_CONFIG_KEY,
+    TIMEOUT_CONFIG_KEY,
+)
+from rasa.shared.core.events import BotUttered, UserUttered
+from rasa.shared.exceptions import AgentInitializationException
+from rasa.shared.providers.llm.llm_response import LLMResponse, LLMToolCall
+from rasa.shared.utils.constants import LOG_COMPONENT_SOURCE_METHOD_INIT
+from rasa.shared.utils.llm import (
+    get_prompt_template,
+    llm_factory,
+    resolve_model_client_config,
+)
+from rasa.shared.utils.mcp.server_connection import MCPServerConnection
+
+DEFAULT_OPENAI_MAX_GENERATED_TOKENS = 256
+MODEL_NAME_GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
+DEFAULT_LLM_CONFIG = {
+    PROVIDER_CONFIG_KEY: OPENAI_PROVIDER,
+    MODEL_CONFIG_KEY: MODEL_NAME_GPT_4O_2024_11_20,
+    TEMPERATURE_CONFIG_KEY: 0.0,
+    MAX_COMPLETION_TOKENS_CONFIG_KEY: DEFAULT_OPENAI_MAX_GENERATED_TOKENS,
+    TIMEOUT_CONFIG_KEY: 7,
+}
+
+
+structlogger = structlog.get_logger()
+
+
+class MCPBaseAgent(AgentProtocol):
+    """MCP protocol implementation"""
+
+    MAX_ITERATIONS = 10
+
+    # ============================================================================
+    # Initialization & Setup
+    # ============================================================================
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        protocol_type: ProtocolConfig,
+        server_configs: List[AgentMCPServerConfig],
+        llm_config: Optional[Dict[str, Any]] = None,
+        prompt_template: Optional[str] = None,
+        timeout: Optional[int] = None,
+        max_retries: Optional[int] = None,
+    ):
+        self._name = name
+
+        self._description = description
+
+        self._protocol_type = protocol_type
+
+        self._llm_config = resolve_model_client_config(
+            llm_config, self.__class__.__name__
+        )
+
+        self.llm_client = llm_factory(self._llm_config, self.get_default_llm_config())
+
+        self.prompt_template = get_prompt_template(
+            prompt_template,
+            self.get_default_prompt_template(),
+            log_source_component=self.__class__.__name__,
+            log_source_method=LOG_COMPONENT_SOURCE_METHOD_INIT,
+        )
+
+        self._timeout = timeout or AGENT_DEFAULT_TIMEOUT_SECONDS
+
+        self._max_retries = max_retries or AGENT_DEFAULT_MAX_RETRIES
+
+        self._server_configs = server_configs or []
+
+        # Stores the MCP tools for the agent.
+        self._mcp_tools: List[AgentToolSchema] = []
+
+        # Stores the custom tools for the agent.
+        self._custom_tools: List[CustomToolSchema] = [
+            CustomToolSchema.from_dict(tool)
+            for tool in self.get_custom_tool_definitions()
+        ]
+
+        # Maps the tool names to the MCP servers that provide them.
+        # key: tool name, value: server name.
+        self._tool_to_server_mapper: Dict[str, str] = {}
+
+        # Stores the connections to the MCP servers.
+        # key: server name, value: connection object.
+        self._server_connections: Dict[str, MCPServerConnection] = {}
+
+    @classmethod
+    def from_config(cls, config: AgentConfig) -> "MCPBaseAgent":
+        """Initialize the MCP Open Agent with the given configuration."""
+        return cls(
+            name=config.agent.name,
+            description=config.agent.description,
+            protocol_type=config.agent.protocol,
+            llm_config=config.configuration.llm if config.configuration else None,
+            prompt_template=config.configuration.prompt_template
+            if config.configuration
+            else None,
+            timeout=config.configuration.timeout if config.configuration else None,
+            max_retries=config.configuration.max_retries
+            if config.configuration
+            else None,
+            server_configs=config.connections.mcp_servers
+            if config.connections
+            else None,
+        )
+
+    # ============================================================================
+    # Class Configuration & Properties
+    # ============================================================================
+
+    @classmethod
+    @abstractmethod
+    def get_default_prompt_template(cls) -> str: ...
+
+    @property
+    def agent_conforms_to(self) -> ProtocolConfig:
+        return self._protocol_type
+
+    @property
+    @abstractmethod
+    def protocol_type(self) -> ProtocolType: ...
+
+    @staticmethod
+    def get_default_llm_config() -> Dict[str, Any]:
+        """Get the default LLM config for the command generator."""
+        return DEFAULT_LLM_CONFIG
+
+    @classmethod
+    def get_agent_specific_built_in_tools(
+        cls, agent_input: AgentInput
+    ) -> List[AgentToolSchema]:
+        """Get agentic specific built-in tools."""
+        return []
+
+    def get_custom_tool_definitions(self) -> List[Dict[str, Any]]:
+        """Add custom tool definitions and their executors for MCP agents.
+
+        This method can be overridden to provide custom tools that the agent
+        can use during its operation. The tools should be defined in the OpenAI
+        JSON format, which includes the tool name, description, and parameters.
+
+        Refer: https://platform.openai.com/docs/guides/function-calling?api-mode=responses
+
+        Note:
+            The tool executor method should be a coroutine function that returns an
+            AgentToolResult object.
+
+        Returns:
+            A list of custom tool definitions along with their tool executor method.
+
+        Example:
+            ```python
+            def get_custom_tool_definitions(self) -> List[Dict[str, Any]]:
+                return [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_current_weather",
+                            "description": "Get the current weather in given location",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "location": {
+                                        "type": "string",
+                                        "description": "The city, e.g. San Francisco",
+                                    },
+                                    "unit": {
+                                        "type": "string",
+                                        "enum": ["celsius", "fahrenheit"],
+                                    },
+                                },
+                                "required": ["location"],
+                            },
+                        },
+                        "tool_executor": self.get_current_weather,
+                    }
+                ]
+            ```
+        """
+        return []
+
+    # ============================================================================
+    # Connection Management
+    # ============================================================================
+
+    async def connect(self) -> None:
+        """Connect to the MCP servers and initialize the agent.
+
+        This method establishes connections to the configured MCP servers,
+        fetches the available tools, and prepares the agent for operation.
+        It should be called before sending any messages to the agent.
+
+        Retries:
+            Retries connection N times if a ConnectionError is raised.
+
+        Logs:
+            Warning: If the connection to any server fails.
+            Warning: If there is a duplicate tool name across servers.
+            Warning: If there is an error fetching tools from any server.
+        """
+        for attempt in range(1, self._max_retries + 1):
+            try:
+                await self.connect_to_servers()
+                await self.fetch_and_store_available_tools()
+                break
+            except ConnectionError as ce:
+                structlogger.warning(
+                    "mcp_agent.connect.connection_error",
+                    event_info=f"Connection attempt {attempt} failed: {ce}",
+                    attempt=attempt,
+                    max_retries=self._max_retries,
+                )
+                if attempt == self._max_retries:
+                    structlogger.error(
+                        "mcp_agent.connect.failed_after_retries",
+                        event_info="All connection attempts failed.",
+                    )
+                    raise AgentInitializationException(
+                        f"Failed to connect to MCP servers after {self._max_retries} "
+                        f"attempts. Agent `{self._name}` failed to initialize."
+                    ) from ce
+            except Exception as e:
+                structlogger.error(
+                    "mcp_agent.connect.unexpected_exception",
+                    event_info="Unexpected error during agent initialization.",
+                    error=str(e),
+                    agent_name=self._name,
+                    agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
+                )
+                raise AgentInitializationException(e) from e
+
+    async def connect_to_server(self, server_config: AgentMCPServerConfig) -> None:
+        server_name = server_config.name
+        connection = MCPServerConnection.from_config(server_config.model_dump())
+        try:
+            await connection.connect()
+            self._server_connections[server_name] = connection
+            structlogger.info(
+                "mcp_agent.connect_to_server.connected",
+                event_info=(
+                    f"Agent `{self._name}` connected to MCP server - "
+                    f"`{server_name}` @ `{connection.server_url}`"
+                ),
+                server_id=server_name,
+                server_url=connection.server_url,
+                agent_name=self._name,
+                agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
+            )
+        except Exception as e:
+            event_info = (
+                f"Agent `{self._name}` failed to connect to MCP server - "
+                f"{server_name} @ {server_config.url}: {e}"
+            )
+            structlogger.error(
+                "mcp_agent.connect.failed_to_connect",
+                event_info=event_info,
+                server_id=server_name,
+                server_url=server_config.url,
+                agent_name=self._name,
+                agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
+            )
+            raise ConnectionError(e)
+
+    async def connect_to_servers(self) -> None:
+        """Connect to MCP servers."""
+        for server_config in self._server_configs:
+            await self.connect_to_server(server_config)
+
+    async def disconnect_server(self, server_name: str) -> None:
+        """Disconnect from an MCP server.
+
+        Args:
+            server_name: The name of the server to disconnect from.
+
+        Side effects:
+            - Removes the server from the server connections.
+
+        Logs:
+            - An error if the server disconnect fails.
+        """
+        if server_name not in self._server_connections:
+            return
+        try:
+            await self._server_connections[server_name].close()
+            del self._server_connections[server_name]
+        except Exception as e:
+            structlogger.error(
+                "mcp_agent.disconnect_server.error",
+                event_info=f"Failed to disconnect from server `{server_name}`: {e}",
+                server_name=server_name,
+                agent_name=self._name,
+                agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
+            )
+
+    async def disconnect(self) -> None:
+        """Close all MCP server connections."""
+        server_names = list(self._server_connections.keys())
+        for server_name in server_names:
+            await self.disconnect_server(server_name)
+
+    # ============================================================================
+    # Tool Management
+    # ============================================================================
+
+    async def list_tools(self, connection: MCPServerConnection) -> ListToolsResult:
+        """List the tools from the MCP server."""
+        session = await connection.ensure_active_session()
+        return await session.list_tools()
+
+    def get_custom_tools(self) -> List[AgentToolSchema]:
+        """Get the custom tools for the agent."""
+        return [tool.tool_definition for tool in self._custom_tools]
+
+    def get_available_tools(self, agent_input: AgentInput) -> List[AgentToolSchema]:
+        """Get the available tools for the agent."""
+        return (
+            self._mcp_tools
+            + self.get_agent_specific_built_in_tools(agent_input)
+            + self.get_custom_tools()
+        )
+
+    async def _get_filtered_tools_from_server(
+        self,
+        server_name: str,
+        connection: MCPServerConnection,
+        include_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+    ) -> List[AgentToolSchema]:
+        """Get filtered tools from the MCP server.
+
+        This method fetches the available tools from the MCP server and filters them
+        based on the include_tools and exclude_tools parameters.
+
+        Args:
+            server_name: The name of the MCP server.
+            connection: The MCP server connection.
+            include_tools: List of tool names to include. If provided, only tools in
+                this list will be fetched.
+            exclude_tools: List of tool names to exclude. If provided, tools in this
+                list will not be fetched.
+
+        Returns:
+            A list of AgentToolSchema objects representing the filtered tools.
+
+        Logs:
+            Warning: If there is a duplicate tool name across servers.
+            Warning: If there is an error fetching tools from the server.
+        """
+        try:
+            tools_response = await self.list_tools(connection)
+            if not tools_response:
+                return []
+
+            filtered_tools = []
+            for tool in tools_response.tools:
+                if include_tools and tool.name not in include_tools:
+                    continue
+                if exclude_tools and tool.name in exclude_tools:
+                    continue
+                filtered_tools.append(AgentToolSchema.from_mcp_tool(tool))
+
+            return filtered_tools
+
+        except Exception as e:
+            event_info = f"Failed to load tools from {server_name}"
+            structlogger.warning(
+                "mcp_agent.get_filtered_tools_from_server.failed_to_get_tools",
+                event_info=event_info,
+                server_name=server_name,
+                server_url=connection.server_url,
+                error=str(e),
+            )
+            return []
+
+    def _get_include_exclude_tools_from_server_configs(
+        self, server_name: str
+    ) -> Tuple[Optional[List[str]], Optional[List[str]]]:
+        """Get the include and exclude tools from the server configs."""
+        for server_config in self._server_configs:
+            if server_config.name == server_name:
+                return server_config.include_tools, server_config.exclude_tools
+        return None, None
+
+    async def fetch_and_store_available_tools(self) -> None:
+        """Fetch and store the available tools from the MCP servers.
+
+        This method fetches the available tools from the MCP servers and stores them
+        in the agent's internal state. It also maps the tool names to the MCP servers
+        that provide them.
+
+        Side effects:
+            - Updates the `_mcp_tools` attribute.
+            - Updates the `_tool_to_server_mapper` attribute.
+
+        Logs:
+            Warning: If there is a duplicate tool name across servers.
+            Warning: If there is an error fetching tools from any server.
+        """
+        for server_name, connection in self._server_connections.items():
+            # Get the include and exclude tools from the server configs.
+            include_tools, exclude_tools = (
+                self._get_include_exclude_tools_from_server_configs(server_name)
+            )
+
+            # Get the filtered tools from the server.
+            tools = await self._get_filtered_tools_from_server(
+                server_name, connection, include_tools, exclude_tools
+            )
+
+            # Add the tools to the tool_to_server_mapper and the available_tools.
+            for tool in tools:
+                if tool.name in self._tool_to_server_mapper:
+                    structlogger.warning(
+                        "mcp_agent.duplicate_tool_name",
+                        event_info=(
+                            f"Tool - {tool.name} from server {server_name} already "
+                            f"exists in {self._tool_to_server_mapper[tool.name]}. "
+                            f"Omitting the tool from server {server_name}."
+                        ),
+                        tool_name=tool.name,
+                        server_name=server_name,
+                        server_url=connection.server_url,
+                    )
+                    continue
+
+                self._tool_to_server_mapper[tool.name] = server_name
+                self._mcp_tools.append(tool)
+
+        structlogger.debug(
+            "mcp_agent.fetch_and_store_available_tools.success",
+            event_info=(
+                "Successfully fetched and stored available tools from MCP servers."
+            ),
+            agent_name=self._name,
+            agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
+            mcp_tools_fetched=len(self._mcp_tools),
+            mcp_tools=[tool.name for tool in self._mcp_tools],
+        )
+
+    # ============================================================================
+    # LLM & Prompt Management
+    # ============================================================================
+
+    def render_prompt_template(self, context: AgentInput) -> str:
+        """Render the prompt template with the provided inputs."""
+        return Template(self.prompt_template).render(
+            **context.model_dump(exclude={"id", "timestamp", "events"}),
+            description=self._description,
+        )
+
+    def build_messages_for_llm_request(
+        self, context: AgentInput, turns: int = 20
+    ) -> List[Dict[str, str]]:
+        """Build messages for the LLM request."""
+        messages = [
+            {KEY_ROLE: ROLE_SYSTEM, KEY_CONTENT: self.render_prompt_template(context)}
+        ]
+
+        # Limit to last N events - set by `turns`.
+        for event in context.events[-turns:]:
+            if isinstance(event, UserUttered):
+                if not event.text:
+                    continue
+                messages.append({KEY_ROLE: ROLE_USER, KEY_CONTENT: event.text})
+            elif isinstance(event, BotUttered):
+                if not event.text:
+                    continue
+                messages.append({KEY_ROLE: ROLE_ASSISTANT, KEY_CONTENT: event.text})
+
+        if context.user_message != messages[-1][KEY_CONTENT]:
+            messages.append({KEY_ROLE: ROLE_USER, KEY_CONTENT: context.user_message})
+        return messages
+
+    def _get_assistant_message_with_tool_calls(
+        self, llm_response: LLMResponse
+    ) -> Dict[str, Any]:
+        """Get assistant message with tool calls."""
+        if not llm_response.tool_calls:
+            return {}
+        return {
+            KEY_ROLE: ROLE_ASSISTANT,
+            KEY_CONTENT: llm_response.choices[0],
+            KEY_TOOL_CALLS: [
+                {
+                    KEY_ID: tool_call.id,
+                    KEY_TYPE: tool_call.type,
+                    KEY_FUNCTION: {
+                        KEY_NAME: tool_call.tool_name,
+                        KEY_ARGUMENTS: json.dumps(tool_call.tool_args),
+                    },
+                }
+                for tool_call in llm_response.tool_calls
+            ],
+        }
+
+    def _get_tool_call_message(self, tool_response: AgentOutput) -> Dict[str, Any]:
+        """Get the tool call message."""
+        return {
+            KEY_ROLE: ROLE_TOOL,
+            KEY_TOOL_CALL_ID: tool_response.id,
+            KEY_CONTENT: tool_response.response_message,
+        }
+
+    # ============================================================================
+    # Tool Execution
+    # ============================================================================
+
+    async def _execute_mcp_tool(
+        self, tool_name: str, arguments: Dict[str, Any]
+    ) -> AgentToolResult:
+        """Execute a tool call via MCP servers."""
+        if tool_name not in self._tool_to_server_mapper:
+            return AgentToolResult(
+                tool_name=tool_name,
+                result=None,
+                is_error=True,
+                error_message=f"Tool `{tool_name}` not found in the server.",
+            )
+
+        server_id = self._tool_to_server_mapper[tool_name]
+        connection = self._server_connections[server_id]
+        session = await connection.ensure_active_session()
+        try:
+            result = await session.call_tool(tool_name, arguments)
+            return AgentToolResult.from_mcp_tool_result(tool_name, result)
+        except Exception as e:
+            return AgentToolResult(
+                tool_name=tool_name,
+                result=None,
+                is_error=True,
+                error_message=(
+                    f"Failed to execute tool `{tool_name}` via MCP server `{server_id}`"
+                    f"@ {connection.server_url}: {e}"
+                ),
+            )
+
+    async def _execute_tool_call(
+        self, tool_name: str, arguments: Dict[str, Any]
+    ) -> AgentToolResult:
+        """Execute a tool call.
+
+        This method first checks if the tool is a built-in tool. If it is, it executes
+        the built-in tool. If it is not, it executes the tool via MCP servers.
+
+        Args:
+            tool_name: The name of the tool to execute.
+            arguments: The arguments to pass to the tool.
+
+        Returns:
+            The result of the tool execution as an AgentToolResult object.
+        """
+        try:
+            for custom_tool in self._custom_tools:
+                if custom_tool.tool_name == tool_name:
+                    result = custom_tool.tool_executor(arguments)
+                    return await result if isawaitable(result) else result
+        except Exception as e:
+            return AgentToolResult(
+                tool_name=tool_name,
+                result=None,
+                is_error=True,
+                error_message=f"Failed to execute built-in tool `{tool_name}`: {e}",
+            )
+        return await self._execute_mcp_tool(tool_name, arguments)
+
+    def _generate_agent_error_output(
+        self,
+        tool_output: AgentToolResult,
+        agent_input: AgentInput,
+        tool_call: LLMToolCall,
+    ) -> AgentOutput:
+        """Generate an agent error output."""
+        structlogger.error(
+            "mcp_agent.send_message.tool_execution_error",
+            event_info=(
+                f"Tool `{tool_output.tool_name}` returned an error: "
+                f"{tool_output.error_message}"
+            ),
+            tool_name=tool_output.tool_name,
+            tool_args=json.dumps(tool_call.tool_args),
+        )
+        if tool_output.is_error:
+            return AgentOutput(
+                id=agent_input.id,
+                status=AgentStatus.FATAL_ERROR,
+                error_message=tool_output.error_message,
+            )
+        else:
+            return AgentOutput(
+                id=agent_input.id,
+                status=AgentStatus.RECOVERABLE_ERROR,
+                error_message=tool_output.error_message,
+            )
+
+    def _get_tool_results_for_agent_output(
+        self,
+        agent_input: AgentInput,
+        current_tool_results: Dict[str, AgentToolResult],
+    ) -> List[List[Dict[str, Any]]]:
+        """Get the tool results for the agent output."""
+        tool_results_of_current_iteration: List[Dict[str, Any]] = []
+        for tool_result in current_tool_results.values():
+            tool_results_of_current_iteration.append(
+                {"tool_name": tool_result.tool_name, "result": tool_result.result}
+            )
+
+        previous_tool_results: List[List[Dict[str, Any]]] = (
+            agent_input.metadata.get(AGENT_METADATA_TOOL_RESULTS_KEY, []) or []
+        )
+        previous_tool_results.append(tool_results_of_current_iteration)
+
+        return previous_tool_results
+
+    # ============================================================================
+    # Core Protocol Methods
+    # ============================================================================
+
+    @abstractmethod
+    async def send_message(self, agent_input: AgentInput) -> AgentOutput:
+        """Send a message to the agent."""
+        ...
+
+    async def run(self, input: AgentInput) -> AgentOutput:
+        """Send a message to Agent/server and return response."""
+        return await self.send_message(input)
+
+    # ============================================================================
+    # Message Processing
+    # ============================================================================
+
+    async def process_input(self, input: AgentInput) -> AgentInput:
+        """Pre-process the input before sending it to the agent."""
+        return input
+
+    async def process_output(self, output: AgentOutput) -> AgentOutput:
+        """Post-process the output before returning it to Rasa."""
+        return output