MCP Service
WIP getting mcp.run to work add mcp[cli] to toml and lint mcp stdio example mcp sse example mcp_run POC ruff formatting
This commit is contained in:
vipyne
213
src/pipecat/services/mcp_service.py
Normal file
213
src/pipecat/services/mcp_service.py
Normal file
@@ -0,0 +1,213 @@
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from pipecat.adapters.schemas.function_schema import FunctionSchema
|
||||
from pipecat.adapters.schemas.tools_schema import ToolsSchema
|
||||
|
||||
try:
|
||||
from mcp import ClientSession, StdioServerParameters, types
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use an MCP client, you need to `pip install pipecat-ai[mcp]`.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class MCPClient:
|
||||
def __init__(
|
||||
self,
|
||||
server_params: Union[StdioServerParameters, str] = None,
|
||||
pipecat_session: Optional[any] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self._server_params = server_params
|
||||
self._session = ClientSession
|
||||
self._pipecat_session = pipecat_session
|
||||
if isinstance(server_params, StdioServerParameters):
|
||||
self._client = stdio_client
|
||||
self._register_tools = self._stdio_register_tools
|
||||
elif isinstance(server_params, str):
|
||||
self._client = sse_client
|
||||
self._register_tools = self._sse_register_tools
|
||||
else:
|
||||
logger.exception(
|
||||
f"{self} error - server_params must be either StdioServerParameters or an sse server url string."
|
||||
)
|
||||
|
||||
async def register_tools(self, llm) -> ToolsSchema:
|
||||
tools_schema = await self._register_tools(llm)
|
||||
return tools_schema
|
||||
|
||||
def convert_mcp_schema_to_pipecat(
|
||||
self, tool_name: str, tool_schema: Dict[str, Any]
|
||||
) -> FunctionSchema:
|
||||
"""Convert an mcp tool schema to Pipecat's FunctionSchema format.
|
||||
Args:
|
||||
tool_name: The name of the tool
|
||||
tool_schema: The mcp tool schema
|
||||
Returns:
|
||||
A FunctionSchema instance
|
||||
"""
|
||||
|
||||
logger.debug(f"Converting schema for tool '{tool_name}'")
|
||||
logger.debug(f"Original schema: {json.dumps(tool_schema, indent=2)}")
|
||||
|
||||
properties = tool_schema["input_schema"].get("properties", {})
|
||||
required = tool_schema["input_schema"].get("required", [])
|
||||
|
||||
schema = FunctionSchema(
|
||||
name=tool_name,
|
||||
description=tool_schema["description"],
|
||||
properties=properties,
|
||||
required=required,
|
||||
)
|
||||
|
||||
logger.debug(f"Converted schema: {json.dumps(schema.to_default_dict(), indent=2)}")
|
||||
|
||||
return schema
|
||||
|
||||
async def _sse_register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available mcp.run tools with the LLM service.
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with
|
||||
Returns:
|
||||
A ToolsSchema containing all registered tools
|
||||
"""
|
||||
|
||||
async def mcp_tool_wrapper(
|
||||
function_name: str,
|
||||
tool_call_id: str,
|
||||
arguments: Dict[str, Any],
|
||||
llm: any,
|
||||
context: any,
|
||||
result_callback: any,
|
||||
) -> None:
|
||||
"""Wrapper for mcp.run tool calls to match Pipecat's function call interface."""
|
||||
logger.debug(f"Executing tool '{function_name}' with call ID: {tool_call_id}")
|
||||
logger.debug(f"Tool arguments: {json.dumps(arguments, indent=2)}")
|
||||
try:
|
||||
async with self._client(self._server_params) as (read, write):
|
||||
async with self._session(read, write) as session:
|
||||
await session.initialize()
|
||||
await self._call_tool(session, function_name, arguments, result_callback)
|
||||
except Exception as e:
|
||||
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):
|
||||
async with self._session(read, write) as session:
|
||||
await session.initialize()
|
||||
tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)
|
||||
return tools_schema
|
||||
|
||||
async def _stdio_register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available mcp.run tools with the LLM service.
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with
|
||||
Returns:
|
||||
A ToolsSchema containing all registered tools
|
||||
"""
|
||||
|
||||
async def mcp_tool_wrapper(
|
||||
function_name: str,
|
||||
tool_call_id: str,
|
||||
arguments: Dict[str, Any],
|
||||
llm: any,
|
||||
context: any,
|
||||
result_callback: any,
|
||||
) -> None:
|
||||
"""Wrapper for mcp.run tool calls to match Pipecat's function call interface."""
|
||||
logger.debug(f"Executing tool '{function_name}' with call ID: {tool_call_id}")
|
||||
logger.debug(f"Tool arguments: {json.dumps(arguments, indent=2)}")
|
||||
try:
|
||||
async with self._client(self._server_params) as streams:
|
||||
async with self._session(streams[0], streams[1]) as session:
|
||||
await session.initialize()
|
||||
await self._call_tool(session, function_name, arguments, result_callback)
|
||||
except Exception as e:
|
||||
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")
|
||||
|
||||
async with self._client(self._server_params) as streams:
|
||||
async with self._session(streams[0], streams[1]) as session:
|
||||
await session.initialize()
|
||||
tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)
|
||||
return tools_schema
|
||||
|
||||
async def _call_tool(self, session, function_name, arguments, result_callback):
|
||||
logger.debug(f"Calling mcp tool '{function_name}'")
|
||||
try:
|
||||
results = await session.call_tool(function_name, arguments=arguments)
|
||||
except Exception as e:
|
||||
error_msg = f"Error calling mcp tool {function_name}: {str(e)}"
|
||||
logger.error(error_msg)
|
||||
|
||||
response = ""
|
||||
image_url = None
|
||||
if hasattr(results, "content") and results.content:
|
||||
for i, content in enumerate(results.content):
|
||||
if hasattr(content, "text") and content.text:
|
||||
logger.debug(f"Tool response chunk {i}: {content.text}")
|
||||
response += content.text
|
||||
else:
|
||||
# logger.debug(f"Non-text result content: '{content}'")
|
||||
pass
|
||||
logger.info(f"Tool '{function_name}' completed successfully")
|
||||
logger.info(f"Final response: {response}")
|
||||
else:
|
||||
logger.error(f"Error getting content from {function_name} results.")
|
||||
|
||||
await result_callback(response)
|
||||
|
||||
async def _list_tools(self, session, mcp_tool_wrapper, llm):
|
||||
available_tools = await session.list_tools()
|
||||
tool_schemas: List[FunctionSchema] = []
|
||||
|
||||
try:
|
||||
logger.debug(f"Found {len(available_tools)} available tools")
|
||||
except:
|
||||
pass
|
||||
|
||||
for tool in available_tools.tools:
|
||||
tool_name = tool.name
|
||||
logger.debug(f"Processing tool: {tool_name}")
|
||||
logger.debug(f"Tool description: {tool.description}")
|
||||
|
||||
try:
|
||||
# Convert the schema
|
||||
function_schema = self.convert_mcp_schema_to_pipecat(
|
||||
tool_name,
|
||||
{"description": tool.description, "input_schema": tool.inputSchema},
|
||||
)
|
||||
|
||||
# Register the wrapped function
|
||||
logger.debug(f"Registering function handler for '{tool_name}'")
|
||||
llm.register_function(tool_name, mcp_tool_wrapper)
|
||||
|
||||
# Add to list of schemas
|
||||
tool_schemas.append(function_schema)
|
||||
logger.debug(f"Successfully registered tool '{tool_name}'")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to register tool '{tool_name}': {str(e)}")
|
||||
logger.exception("Full exception details:")
|
||||
continue
|
||||
|
||||
logger.info(f"Completed registration of {len(tool_schemas)} tools")
|
||||
tools_schema = ToolsSchema(standard_tools=tool_schemas)
|
||||
|
||||
return tools_schema
|
||||
Reference in New Issue
Block a user