Source code for autogen_ext.tools.mcp._host._session_host

from typing import Any, Dict

from autogen_core import Component, ComponentBase, ComponentModel
from pydantic import BaseModel

from mcp import types as mcp_types

from ._elicitation import Elicitor
from ._roots import RootsProvider
from ._sampling import Sampler


class McpSessionHostConfig(BaseModel):
    """Configuration for MCP session host components.

    Args:
        model_client: Optional chat completion client for sampling requests
        elicitor: Optional elicitor component for handling elicitation requests
        roots: Optional list of file system roots or roots provider
    """

    sampler: ComponentModel | Dict[str, Any] | None
    elicitor: ComponentModel | Dict[str, Any] | None
    roots: ComponentModel | Dict[str, Any] | None


[docs] class McpSessionHost(ComponentBase[BaseModel], Component[McpSessionHostConfig]): """Host component that provides MCP server capabilities. This host acts as the client-side Host for MCP sessions, handling requests from MCP servers for text generation (sampling), user prompting (elicitation), and file system root listing. It coordinates with model clients and elicitors to provide these capabilities. The host supports three main MCP server capabilities: - Sampling: Text generation using a language model via chat completion client - Elicitation: Structured prompting and response collection via elicitors - Roots: Listing available file system roots for server access Args: model_client: Optional chat completion client for handling sampling requests roots: Optional sequence of roots or callable returning roots for file system access elicitor: Optional elicitor for handling elicitation requests Example: Complete setup with MCP capabilities including sampling and elicitation:: from autogen_agentchat.agents import AssistantAgent from autogen_ext.models.openai import OpenAIChatCompletionClient from autogen_ext.tools.mcp import ( ChatCompletionClientSampler, McpSessionHost, McpWorkbench, StaticRootsProvider, StdioElicitor, StdioServerParams, ) from pydantic import FileUrl from mcp.types import Root # Setup model client for sampling model_client = OpenAIChatCompletionClient(model="gpt-4o") sampler = ChatCompletionClientSampler(model_client) # Create elicitor that prompts for user input over stdio elicitor = StdioElicitor() # Provide static roots in the host system roots = StaticRootsProvider( [Root(uri=FileUrl("file:///home"), name="Home"), Root(uri=FileUrl("file:///tmp"), name="Tmp")] ) # Create MCP session host with sampling, elicitation, and list_roots capabilities # If you want to support roots, import or define Root and FileUrl, then uncomment the roots line below host = McpSessionHost( sampler=sampler, # Support sampling via model client elicitor=elicitor, # Support elicitation via user_proxy roots=roots, ) # Setup MCP workbench with your server mcp_workbench = McpWorkbench( server_params=StdioServerParams(command="python", args=["your_mcp_server.py"]), host=host, # Add the host here ) # Create MCP-enabled assistant mcp_assistant = AssistantAgent( "mcp_assistant", model_client=model_client, workbench=mcp_workbench, ) # Now the AssistantAgent can support MCP servers that request sampling, elicitation, and roots! """ component_type = "mcp_session_host" component_config_schema = McpSessionHostConfig component_provider_override = "autogen_ext.tools.mcp.McpSessionHost" def __init__( self, sampler: Sampler | None = None, roots: RootsProvider | None = None, elicitor: Elicitor | None = None, ): """Initialize the MCP session host. Args: sampler: Optional sampler handling sampling requests. roots: Optional roots provider for returning roots for file system access. elicitor: Optional elicitor for handling elicitation requests. """ self._sampler = sampler self._roots = roots self._elicitor = elicitor
[docs] async def handle_sampling_request( self, params: mcp_types.CreateMessageRequestParams ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData: """Handle a sampling request from MCP servers. Converts MCP messages to AutoGen format and uses the configured sampler (if any) to generate a response. Args: params: The sampling request containing message creation parameters. Returns: A sampling response with the generated message or error data. """ if self._sampler is None: return mcp_types.ErrorData( code=mcp_types.INVALID_REQUEST, message="No model client available for sampling requests", ) try: response = await self._sampler.sample(params) return response except Exception as e: return mcp_types.ErrorData( code=mcp_types.INTERNAL_ERROR, message=f"Sampling request failed: {str(e)}", )
[docs] async def handle_elicit_request( self, params: mcp_types.ElicitRequestParams ) -> mcp_types.ElicitResult | mcp_types.ErrorData: """Handle an elicitation request from MCP servers. Forwards the elicitation request to the configured elicitor for processing. The elicitor handles the structured prompting and response collection. Args: params: The elicitation request containing prompts and parameters. Returns: An elicitation response with the structured result or error data. """ if self._elicitor is None: return mcp_types.ErrorData( code=mcp_types.INVALID_REQUEST, message="No elicitor configured for this host", ) try: return await self._elicitor.elicit(params) except Exception as e: return mcp_types.ErrorData( code=mcp_types.INTERNAL_ERROR, message=f"Elicitation request failed: {str(e)}", )
[docs] async def handle_list_roots_request(self) -> mcp_types.ListRootsResult | mcp_types.ErrorData: """Handle a list roots request from MCP servers. Returns the configured file system roots that are available for server access. Returns: A list roots response containing available roots or error data. """ if self._roots is None: return mcp_types.ErrorData(code=mcp_types.INVALID_REQUEST, message="Host does not support listing roots") else: try: return await self._roots.list_roots() except Exception as e: return mcp_types.ErrorData(code=mcp_types.INTERNAL_ERROR, message=f"Caught error listing roots: {e}")
[docs] def _to_config(self) -> BaseModel: return McpSessionHostConfig( sampler=self._sampler.dump_component() if self._sampler else None, elicitor=self._elicitor.dump_component() if self._elicitor else None, roots=self._roots.dump_component() if self._roots else None, )
[docs] @classmethod def _from_config(cls, config: McpSessionHostConfig) -> "McpSessionHost": return cls( sampler=Sampler.load_component(config.sampler) if config.sampler else None, elicitor=Elicitor.load_component(config.elicitor) if config.elicitor else None, roots=RootsProvider.load_component(config.roots) if config.roots else None, )