diff --git a/src/pipecat/services/mcp_service.py b/src/pipecat/services/mcp_service.py index ae54b84ff..2e519dd24 100644 --- a/src/pipecat/services/mcp_service.py +++ b/src/pipecat/services/mcp_service.py @@ -8,8 +8,8 @@ from pipecat.adapters.schemas.tools_schema import ToolsSchema from pipecat.utils.base_object import BaseObject try: - from mcp import ClientSession, StdioServerParameters, types - from mcp.client.session import ClientSession + from mcp import ClientSession, StdioServerParameters + from mcp.client.session_group import SseServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client except ModuleNotFoundError as e: @@ -21,7 +21,7 @@ except ModuleNotFoundError as e: class MCPClient(BaseObject): def __init__( self, - server_params: Union[StdioServerParameters, str], + server_params: Union[StdioServerParameters, SseServerParameters], **kwargs, ): super().__init__(**kwargs) @@ -30,12 +30,12 @@ class MCPClient(BaseObject): if isinstance(server_params, StdioServerParameters): self._client = stdio_client self._register_tools = self._stdio_register_tools - elif isinstance(server_params, str): + elif isinstance(server_params, SseServerParameters): self._client = sse_client self._register_tools = self._sse_register_tools else: raise TypeError( - f"{self} invalid argument type: `server_params` must be either StdioServerParameters or an SSE server url string." + f"{self} invalid argument type: `server_params` must be either StdioServerParameters or SseServerParameters." ) async def register_tools(self, llm) -> ToolsSchema: @@ -90,7 +90,12 @@ class MCPClient(BaseObject): logger.debug(f"Executing tool '{function_name}' with call ID: {tool_call_id}") logger.trace(f"Tool arguments: {json.dumps(arguments, indent=2)}") try: - async with self._client(self._server_params) as (read, write): + async with self._client( + url=self._server_params.url, + headers=self._server_params.headers, + timeout=self._server_params.timeout, + sse_read_timeout=self._server_params.sse_read_timeout + ) as (read, write): async with self._session(read, write) as session: await session.initialize() await self._call_tool(session, function_name, arguments, result_callback) @@ -98,12 +103,16 @@ class MCPClient(BaseObject): error_msg = f"Error calling mcp tool {function_name}: {str(e)}" logger.error(error_msg) logger.exception("Full exception details:") - await result_callback(error_msg) - - logger.debug("Starting registration of mcp.run tools") - tool_schemas: List[FunctionSchema] = [] - - async with self._client(self._server_params) as (read, write): + await result_callback(error_msg)\ + + logger.debug(f"SSE server parameters: {self._server_params}") + + async with self._client( + url=self._server_params.url, + headers=self._server_params.headers, + timeout=self._server_params.timeout, + sse_read_timeout=self._server_params.sse_read_timeout + ) as (read, write): async with self._session(read, write) as session: await session.initialize() tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)