Add MCPClient docstrings. Removed google specific cleanup, changed example to openai
This commit is contained in:
committed by
Yousif Astarabadi
parent
1cac028bfe
commit
f0bcc9d9ba
@@ -63,6 +63,10 @@ async def run_example(transport: BaseTransport, _: argparse.Namespace, handle_si
|
||||
)
|
||||
|
||||
try:
|
||||
# Github MCP docs: https://github.com/github/github-mcp-server
|
||||
# Enable Github Copilot on your GitHub account. Free tier is ok. (https://github.com/settings/copilot)
|
||||
# Generate a personal access token. It must be a Fine-grained token, classic tokens are not supported. (https://github.com/settings/personal-access-tokens)
|
||||
# Set permissions you want to use (eg. "all repositories", "profile: read/write", etc)
|
||||
mcp = MCPClient(
|
||||
server_params=StreamableHttpParameters(
|
||||
url="https://api.githubcopilot.com/mcp/",
|
||||
@@ -128,4 +132,4 @@ async def run_example(transport: BaseTransport, _: argparse.Namespace, handle_si
|
||||
if __name__ == "__main__":
|
||||
from pipecat.examples.run import main
|
||||
|
||||
main(run_example, transport_params=transport_params)
|
||||
main(run_example, transport_params=transport_params)
|
||||
@@ -1,5 +1,13 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""MCP (Model Context Protocol) client for integrating external tools with LLMs."""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -10,6 +18,7 @@ from pipecat.utils.base_object import BaseObject
|
||||
try:
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.session_group import SseServerParameters, StreamableHttpParameters
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
@@ -20,15 +29,29 @@ except ModuleNotFoundError as e:
|
||||
|
||||
|
||||
class MCPClient(BaseObject):
|
||||
"""Client for Model Context Protocol (MCP) servers.
|
||||
|
||||
Enables integration with MCP servers to provide external tools and resources
|
||||
to LLMs. Supports both stdio and SSE server connections with automatic tool
|
||||
registration and schema conversion.
|
||||
|
||||
Args:
|
||||
server_params: Server connection parameters (stdio or SSE).
|
||||
**kwargs: Additional arguments passed to the parent BaseObject.
|
||||
|
||||
Raises:
|
||||
TypeError: If server_params is not a supported parameter type.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_params: Union[StdioServerParameters, SseServerParameters, StreamableHttpParameters],
|
||||
server_params: Tuple[StdioServerParameters, SseServerParameters, StreamableHttpParameters],
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self._server_params = server_params
|
||||
self._session = ClientSession
|
||||
|
||||
|
||||
if isinstance(server_params, StdioServerParameters):
|
||||
self._client = stdio_client
|
||||
self._register_tools = self._stdio_register_tools
|
||||
@@ -44,20 +67,32 @@ class MCPClient(BaseObject):
|
||||
)
|
||||
|
||||
async def register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available MCP tools with an LLM service.
|
||||
|
||||
Connects to the MCP server, discovers available tools, converts their
|
||||
schemas to Pipecat format, and registers them with the LLM service.
|
||||
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with.
|
||||
|
||||
Returns:
|
||||
A ToolsSchema containing all successfully registered tools.
|
||||
"""
|
||||
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.trace(f"Original schema: {json.dumps(tool_schema, indent=2)}")
|
||||
|
||||
@@ -76,7 +111,8 @@ class MCPClient(BaseObject):
|
||||
return schema
|
||||
|
||||
async def _sse_register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available mcp.run tools with the LLM service.
|
||||
"""Register all available mcp tools with the LLM service.
|
||||
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with
|
||||
Returns:
|
||||
@@ -91,15 +127,12 @@ class MCPClient(BaseObject):
|
||||
context: any,
|
||||
result_callback: any,
|
||||
) -> None:
|
||||
"""Wrapper for mcp.run tool calls to match Pipecat's function call interface."""
|
||||
"""Wrapper for mcp tool calls to match Pipecat's function call interface."""
|
||||
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(
|
||||
url=self._server_params.url,
|
||||
headers=self._server_params.headers,
|
||||
timeout=self._server_params.timeout,
|
||||
sse_read_timeout=self._server_params.sse_read_timeout,
|
||||
**self._server_params.model_dump()
|
||||
) as (read, write):
|
||||
async with self._session(read, write) as session:
|
||||
await session.initialize()
|
||||
@@ -111,12 +144,10 @@ class MCPClient(BaseObject):
|
||||
await result_callback(error_msg)
|
||||
|
||||
logger.debug(f"SSE server parameters: {self._server_params}")
|
||||
logger.debug("Starting registration of mcp tools")
|
||||
|
||||
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,
|
||||
**self._server_params.model_dump()
|
||||
) as (read, write):
|
||||
async with self._session(read, write) as session:
|
||||
await session.initialize()
|
||||
@@ -124,7 +155,8 @@ class MCPClient(BaseObject):
|
||||
return tools_schema
|
||||
|
||||
async def _stdio_register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available mcp.run tools with the LLM service.
|
||||
"""Register all available mcp tools with the LLM service.
|
||||
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with
|
||||
Returns:
|
||||
@@ -139,7 +171,7 @@ class MCPClient(BaseObject):
|
||||
context: any,
|
||||
result_callback: any,
|
||||
) -> None:
|
||||
"""Wrapper for mcp.run tool calls to match Pipecat's function call interface."""
|
||||
"""Wrapper for mcp tool calls to match Pipecat's function call interface."""
|
||||
logger.debug(f"Executing tool '{function_name}' with call ID: {tool_call_id}")
|
||||
logger.trace(f"Tool arguments: {json.dumps(arguments, indent=2)}")
|
||||
try:
|
||||
@@ -153,7 +185,7 @@ class MCPClient(BaseObject):
|
||||
logger.exception("Full exception details:")
|
||||
await result_callback(error_msg)
|
||||
|
||||
logger.debug("Starting registration of mcp.run tools")
|
||||
logger.debug("Starting registration of mcp tools")
|
||||
|
||||
async with self._client(self._server_params) as streams:
|
||||
async with self._session(streams[0], streams[1]) as session:
|
||||
@@ -162,7 +194,7 @@ class MCPClient(BaseObject):
|
||||
return tools_schema
|
||||
|
||||
async def _streamable_http_register_tools(self, llm) -> ToolsSchema:
|
||||
"""Register all available mcp.run tools with the LLM service using streamable HTTP.
|
||||
"""Register all available mcp tools with the LLM service using streamable HTTP.
|
||||
Args:
|
||||
llm: The Pipecat LLM service to register tools with
|
||||
Returns:
|
||||
@@ -177,16 +209,17 @@ class MCPClient(BaseObject):
|
||||
context: any,
|
||||
result_callback: any,
|
||||
) -> None:
|
||||
"""Wrapper for mcp.run tool calls to match Pipecat's function call interface."""
|
||||
"""Wrapper for mcp tool calls to match Pipecat's function call interface."""
|
||||
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(
|
||||
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_stream, write_stream, _):
|
||||
**self._server_params.model_dump()
|
||||
) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with self._session(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
await self._call_tool(session, function_name, arguments, result_callback)
|
||||
@@ -196,14 +229,15 @@ class MCPClient(BaseObject):
|
||||
logger.exception("Full exception details:")
|
||||
await result_callback(error_msg)
|
||||
|
||||
logger.debug("Starting registration of mcp.run tools using streamable HTTP")
|
||||
logger.debug("Starting registration of mcp tools using streamable HTTP")
|
||||
|
||||
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_stream, write_stream, _):
|
||||
**self._server_params.model_dump()
|
||||
) as (
|
||||
read_stream,
|
||||
write_stream,
|
||||
_,
|
||||
):
|
||||
async with self._session(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)
|
||||
@@ -253,7 +287,7 @@ class MCPClient(BaseObject):
|
||||
# Convert the schema
|
||||
function_schema = self._convert_mcp_schema_to_pipecat(
|
||||
tool_name,
|
||||
{"description": tool.description, "input_schema": tool.inputSchema},
|
||||
{"description": tool.description, "input_schema": tool.inputSchema}
|
||||
)
|
||||
|
||||
# Register the wrapped function
|
||||
@@ -272,4 +306,4 @@ class MCPClient(BaseObject):
|
||||
logger.debug(f"Completed registration of {len(tool_schemas)} tools")
|
||||
tools_schema = ToolsSchema(standard_tools=tool_schemas)
|
||||
|
||||
return tools_schema
|
||||
return tools_schema
|
||||
Reference in New Issue
Block a user