Kye Gomez
3 weeks ago
parent
9cb2500e58
commit
89fc8c7609
@ -0,0 +1,554 @@
|
||||
from contextlib import AsyncExitStack
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Coroutine,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
TypedDict,
|
||||
cast,
|
||||
)
|
||||
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
EmbeddedResource,
|
||||
ImageContent,
|
||||
PromptMessage,
|
||||
TextContent,
|
||||
)
|
||||
from mcp.types import (
|
||||
Tool as MCPTool,
|
||||
)
|
||||
|
||||
|
||||
def convert_mcp_prompt_message_to_message(
|
||||
message: PromptMessage,
|
||||
) -> str:
|
||||
"""Convert an MCP prompt message to a string message.
|
||||
|
||||
Args:
|
||||
message: MCP prompt message to convert
|
||||
|
||||
Returns:
|
||||
a string message
|
||||
"""
|
||||
if message.content.type == "text":
|
||||
if message.role == "user":
|
||||
return str(message.content.text)
|
||||
elif message.role == "assistant":
|
||||
return str(
|
||||
message.content.text
|
||||
) # Fixed attribute name from str to text
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported prompt message role: {message.role}"
|
||||
)
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported prompt message content type: {message.content.type}"
|
||||
)
|
||||
|
||||
|
||||
async def load_mcp_prompt(
|
||||
session: ClientSession,
|
||||
name: str,
|
||||
arguments: Optional[dict[str, Any]] = None,
|
||||
) -> List[str]:
|
||||
"""Load MCP prompt and convert to messages."""
|
||||
response = await session.get_prompt(name, arguments)
|
||||
|
||||
return [
|
||||
convert_mcp_prompt_message_to_message(message)
|
||||
for message in response.messages
|
||||
]
|
||||
|
||||
|
||||
DEFAULT_ENCODING = "utf-8"
|
||||
DEFAULT_ENCODING_ERROR_HANDLER = "strict"
|
||||
|
||||
DEFAULT_HTTP_TIMEOUT = 5
|
||||
DEFAULT_SSE_READ_TIMEOUT = 60 * 5
|
||||
|
||||
|
||||
class StdioConnection(TypedDict):
|
||||
transport: Literal["stdio"]
|
||||
|
||||
command: str
|
||||
"""The executable to run to start the server."""
|
||||
|
||||
args: list[str]
|
||||
"""Command line arguments to pass to the executable."""
|
||||
|
||||
env: dict[str, str] | None
|
||||
"""The environment to use when spawning the process."""
|
||||
|
||||
encoding: str
|
||||
"""The text encoding used when sending/receiving messages to the server."""
|
||||
|
||||
encoding_error_handler: Literal["strict", "ignore", "replace"]
|
||||
"""
|
||||
The text encoding error handler.
|
||||
|
||||
See https://docs.python.org/3/library/codecs.html#codec-base-classes for
|
||||
explanations of possible values
|
||||
"""
|
||||
|
||||
|
||||
class SSEConnection(TypedDict):
|
||||
transport: Literal["sse"]
|
||||
|
||||
url: str
|
||||
"""The URL of the SSE endpoint to connect to."""
|
||||
|
||||
headers: dict[str, Any] | None
|
||||
"""HTTP headers to send to the SSE endpoint"""
|
||||
|
||||
timeout: float
|
||||
"""HTTP timeout"""
|
||||
|
||||
sse_read_timeout: float
|
||||
"""SSE read timeout"""
|
||||
|
||||
|
||||
NonTextContent = ImageContent | EmbeddedResource
|
||||
|
||||
|
||||
def _convert_call_tool_result(
|
||||
call_tool_result: CallToolResult,
|
||||
) -> tuple[str | list[str], list[NonTextContent] | None]:
|
||||
text_contents: list[TextContent] = []
|
||||
non_text_contents = []
|
||||
for content in call_tool_result.content:
|
||||
if isinstance(content, TextContent):
|
||||
text_contents.append(content)
|
||||
else:
|
||||
non_text_contents.append(content)
|
||||
|
||||
tool_content: str | list[str] = [
|
||||
content.text for content in text_contents
|
||||
]
|
||||
if len(text_contents) == 1:
|
||||
tool_content = tool_content[0]
|
||||
|
||||
if call_tool_result.isError:
|
||||
raise ValueError("Error calling tool")
|
||||
|
||||
return tool_content, non_text_contents or None
|
||||
|
||||
|
||||
def convert_mcp_tool_to_function(
|
||||
session: ClientSession,
|
||||
tool: MCPTool,
|
||||
) -> Callable[
|
||||
...,
|
||||
Coroutine[
|
||||
Any, Any, tuple[str | list[str], list[NonTextContent] | None]
|
||||
],
|
||||
]:
|
||||
"""Convert an MCP tool to a callable function.
|
||||
|
||||
NOTE: this tool can be executed only in a context of an active MCP client session.
|
||||
|
||||
Args:
|
||||
session: MCP client session
|
||||
tool: MCP tool to convert
|
||||
|
||||
Returns:
|
||||
a callable function
|
||||
"""
|
||||
|
||||
async def call_tool(
|
||||
**arguments: dict[str, Any],
|
||||
) -> tuple[str | list[str], list[NonTextContent] | None]:
|
||||
"""Execute the tool with the given arguments."""
|
||||
call_tool_result = await session.call_tool(
|
||||
tool.name, arguments
|
||||
)
|
||||
return _convert_call_tool_result(call_tool_result)
|
||||
|
||||
# Add metadata as attributes to the function
|
||||
call_tool.__name__ = tool.name
|
||||
call_tool.__doc__ = tool.description or ""
|
||||
call_tool.schema = tool.inputSchema
|
||||
|
||||
return call_tool
|
||||
|
||||
|
||||
async def load_mcp_tools(session: ClientSession) -> list[Callable]:
|
||||
"""Load all available MCP tools and convert them to callable functions."""
|
||||
tools = await session.list_tools()
|
||||
return [
|
||||
convert_mcp_tool_to_function(session, tool)
|
||||
for tool in tools.tools
|
||||
]
|
||||
|
||||
|
||||
class MultiServerMCPClient:
|
||||
"""Client for connecting to multiple MCP servers and loading tools from them."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connections: dict[
|
||||
str, StdioConnection | SSEConnection
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Initialize a MultiServerMCPClient with MCP servers connections.
|
||||
|
||||
Args:
|
||||
connections: A dictionary mapping server names to connection configurations.
|
||||
Each configuration can be either a StdioConnection or SSEConnection.
|
||||
If None, no initial connections are established.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
async with MultiServerMCPClient(
|
||||
{
|
||||
"math": {
|
||||
"command": "python",
|
||||
# Make sure to update to the full absolute path to your math_server.py file
|
||||
"args": ["/path/to/math_server.py"],
|
||||
"transport": "stdio",
|
||||
},
|
||||
"weather": {
|
||||
# make sure you start your weather server on port 8000
|
||||
"url": "http://localhost:8000/sse",
|
||||
"transport": "sse",
|
||||
}
|
||||
}
|
||||
) as client:
|
||||
all_tools = client.get_tools()
|
||||
...
|
||||
```
|
||||
"""
|
||||
self.connections = connections
|
||||
self.exit_stack = AsyncExitStack()
|
||||
self.sessions: dict[str, ClientSession] = {}
|
||||
self.server_name_to_tools: dict[str, list[Callable]] = {}
|
||||
|
||||
async def _initialize_session_and_load_tools(
|
||||
self, server_name: str, session: ClientSession
|
||||
) -> None:
|
||||
"""Initialize a session and load tools from it.
|
||||
|
||||
Args:
|
||||
server_name: Name to identify this server connection
|
||||
session: The ClientSession to initialize
|
||||
"""
|
||||
# Initialize the session
|
||||
await session.initialize()
|
||||
self.sessions[server_name] = session
|
||||
|
||||
# Load tools from this server
|
||||
server_tools = await load_mcp_tools(session)
|
||||
self.server_name_to_tools[server_name] = server_tools
|
||||
|
||||
async def connect_to_server(
|
||||
self,
|
||||
server_name: str,
|
||||
*,
|
||||
transport: Literal["stdio", "sse"] = "stdio",
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""Connect to an MCP server using either stdio or SSE.
|
||||
|
||||
This is a generic method that calls either connect_to_server_via_stdio or connect_to_server_via_sse
|
||||
based on the provided transport parameter.
|
||||
|
||||
Args:
|
||||
server_name: Name to identify this server connection
|
||||
transport: Type of transport to use ("stdio" or "sse"), defaults to "stdio"
|
||||
**kwargs: Additional arguments to pass to the specific connection method
|
||||
|
||||
Raises:
|
||||
ValueError: If transport is not recognized
|
||||
ValueError: If required parameters for the specified transport are missing
|
||||
"""
|
||||
if transport == "sse":
|
||||
if "url" not in kwargs:
|
||||
raise ValueError(
|
||||
"'url' parameter is required for SSE connection"
|
||||
)
|
||||
await self.connect_to_server_via_sse(
|
||||
server_name,
|
||||
url=kwargs["url"],
|
||||
headers=kwargs.get("headers"),
|
||||
timeout=kwargs.get("timeout", DEFAULT_HTTP_TIMEOUT),
|
||||
sse_read_timeout=kwargs.get(
|
||||
"sse_read_timeout", DEFAULT_SSE_READ_TIMEOUT
|
||||
),
|
||||
)
|
||||
elif transport == "stdio":
|
||||
if "command" not in kwargs:
|
||||
raise ValueError(
|
||||
"'command' parameter is required for stdio connection"
|
||||
)
|
||||
if "args" not in kwargs:
|
||||
raise ValueError(
|
||||
"'args' parameter is required for stdio connection"
|
||||
)
|
||||
await self.connect_to_server_via_stdio(
|
||||
server_name,
|
||||
command=kwargs["command"],
|
||||
args=kwargs["args"],
|
||||
env=kwargs.get("env"),
|
||||
encoding=kwargs.get("encoding", DEFAULT_ENCODING),
|
||||
encoding_error_handler=kwargs.get(
|
||||
"encoding_error_handler",
|
||||
DEFAULT_ENCODING_ERROR_HANDLER,
|
||||
),
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported transport: {transport}. Must be 'stdio' or 'sse'"
|
||||
)
|
||||
|
||||
async def connect_to_server_via_stdio(
|
||||
self,
|
||||
server_name: str,
|
||||
*,
|
||||
command: str,
|
||||
args: list[str],
|
||||
env: dict[str, str] | None = None,
|
||||
encoding: str = DEFAULT_ENCODING,
|
||||
encoding_error_handler: Literal[
|
||||
"strict", "ignore", "replace"
|
||||
] = DEFAULT_ENCODING_ERROR_HANDLER,
|
||||
) -> None:
|
||||
"""Connect to a specific MCP server using stdio
|
||||
|
||||
Args:
|
||||
server_name: Name to identify this server connection
|
||||
command: Command to execute
|
||||
args: Arguments for the command
|
||||
env: Environment variables for the command
|
||||
encoding: Character encoding
|
||||
encoding_error_handler: How to handle encoding errors
|
||||
"""
|
||||
server_params = StdioServerParameters(
|
||||
command=command,
|
||||
args=args,
|
||||
env=env,
|
||||
encoding=encoding,
|
||||
encoding_error_handler=encoding_error_handler,
|
||||
)
|
||||
|
||||
# Create and store the connection
|
||||
stdio_transport = await self.exit_stack.enter_async_context(
|
||||
stdio_client(server_params)
|
||||
)
|
||||
read, write = stdio_transport
|
||||
session = cast(
|
||||
ClientSession,
|
||||
await self.exit_stack.enter_async_context(
|
||||
ClientSession(read, write)
|
||||
),
|
||||
)
|
||||
|
||||
await self._initialize_session_and_load_tools(
|
||||
server_name, session
|
||||
)
|
||||
|
||||
async def connect_to_server_via_sse(
|
||||
self,
|
||||
server_name: str,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, Any] | None = None,
|
||||
timeout: float = DEFAULT_HTTP_TIMEOUT,
|
||||
sse_read_timeout: float = DEFAULT_SSE_READ_TIMEOUT,
|
||||
) -> None:
|
||||
"""Connect to a specific MCP server using SSE
|
||||
|
||||
Args:
|
||||
server_name: Name to identify this server connection
|
||||
url: URL of the SSE server
|
||||
headers: HTTP headers to send to the SSE endpoint
|
||||
timeout: HTTP timeout
|
||||
sse_read_timeout: SSE read timeout
|
||||
"""
|
||||
# Create and store the connection
|
||||
sse_transport = await self.exit_stack.enter_async_context(
|
||||
sse_client(url, headers, timeout, sse_read_timeout)
|
||||
)
|
||||
read, write = sse_transport
|
||||
session = cast(
|
||||
ClientSession,
|
||||
await self.exit_stack.enter_async_context(
|
||||
ClientSession(read, write)
|
||||
),
|
||||
)
|
||||
|
||||
await self._initialize_session_and_load_tools(
|
||||
server_name, session
|
||||
)
|
||||
|
||||
def get_tools(self) -> list[Callable]:
|
||||
"""Get a list of all tools from all connected servers."""
|
||||
all_tools: list[Callable] = []
|
||||
for server_tools in self.server_name_to_tools.values():
|
||||
all_tools.extend(server_tools)
|
||||
return all_tools
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
server_name: str,
|
||||
prompt_name: str,
|
||||
arguments: Optional[dict[str, Any]] = None,
|
||||
) -> List[str]:
|
||||
"""Get a prompt from a given MCP server."""
|
||||
session = self.sessions[server_name]
|
||||
return await load_mcp_prompt(session, prompt_name, arguments)
|
||||
|
||||
async def __aenter__(self) -> "MultiServerMCPClient":
|
||||
try:
|
||||
connections = self.connections or {}
|
||||
for server_name, connection in connections.items():
|
||||
connection_dict = connection.copy()
|
||||
transport = connection_dict.pop("transport")
|
||||
if transport == "stdio":
|
||||
await self.connect_to_server_via_stdio(
|
||||
server_name, **connection_dict
|
||||
)
|
||||
elif transport == "sse":
|
||||
await self.connect_to_server_via_sse(
|
||||
server_name, **connection_dict
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported transport: {transport}. Must be 'stdio' or 'sse'"
|
||||
)
|
||||
return self
|
||||
except Exception:
|
||||
await self.exit_stack.aclose()
|
||||
raise
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
await self.exit_stack.aclose()
|
||||
|
||||
|
||||
#!/usr/bin/env python3
|
||||
import asyncio
|
||||
import os
|
||||
import json
|
||||
from typing import List, Any, Callable
|
||||
|
||||
# # Import our MCP client module
|
||||
# from mcp_client import MultiServerMCPClient
|
||||
|
||||
|
||||
async def main():
|
||||
"""Test script for demonstrating MCP client usage."""
|
||||
print("Starting MCP Client test...")
|
||||
|
||||
# Create a connection to multiple MCP servers
|
||||
# You'll need to update these paths to match your setup
|
||||
async with MultiServerMCPClient(
|
||||
{
|
||||
"math": {
|
||||
"transport": "stdio",
|
||||
"command": "python",
|
||||
"args": ["/path/to/math_server.py"],
|
||||
"env": {"DEBUG": "1"},
|
||||
},
|
||||
"search": {
|
||||
"transport": "sse",
|
||||
"url": "http://localhost:8000/sse",
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {os.environ.get('API_KEY', '')}"
|
||||
},
|
||||
},
|
||||
}
|
||||
) as client:
|
||||
# Get all available tools
|
||||
tools = client.get_tools()
|
||||
print(f"Found {len(tools)} tools across all servers")
|
||||
|
||||
# Print tool information
|
||||
for i, tool in enumerate(tools):
|
||||
print(f"\nTool {i+1}: {tool.__name__}")
|
||||
print(f" Description: {tool.__doc__}")
|
||||
if hasattr(tool, "schema") and tool.schema:
|
||||
print(
|
||||
f" Schema: {json.dumps(tool.schema, indent=2)[:100]}..."
|
||||
)
|
||||
|
||||
# Example: Use a specific tool if available
|
||||
calculator_tool = next(
|
||||
(t for t in tools if t.__name__ == "calculator"), None
|
||||
)
|
||||
if calculator_tool:
|
||||
print("\n\nTesting calculator tool:")
|
||||
try:
|
||||
# Call the tool as an async function
|
||||
result, artifacts = await calculator_tool(
|
||||
expression="2 + 2 * 3"
|
||||
)
|
||||
print(f" Calculator result: {result}")
|
||||
if artifacts:
|
||||
print(
|
||||
f" With {len(artifacts)} additional artifacts"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f" Error using calculator: {e}")
|
||||
|
||||
# Example: Load a prompt from a server
|
||||
try:
|
||||
print("\n\nTesting prompt loading:")
|
||||
prompt_messages = await client.get_prompt(
|
||||
"math",
|
||||
"calculation_introduction",
|
||||
{"user_name": "Test User"},
|
||||
)
|
||||
print(
|
||||
f" Loaded prompt with {len(prompt_messages)} messages:"
|
||||
)
|
||||
for i, msg in enumerate(prompt_messages):
|
||||
print(f" Message {i+1}: {msg[:50]}...")
|
||||
except Exception as e:
|
||||
print(f" Error loading prompt: {e}")
|
||||
|
||||
|
||||
async def create_custom_tool():
|
||||
"""Example of creating a custom tool function."""
|
||||
|
||||
# Define a tool function with metadata
|
||||
async def add_numbers(a: float, b: float) -> tuple[str, None]:
|
||||
"""Add two numbers together."""
|
||||
result = a + b
|
||||
return f"The sum of {a} and {b} is {result}", None
|
||||
|
||||
# Add metadata to the function
|
||||
add_numbers.__name__ = "add_numbers"
|
||||
add_numbers.__doc__ = (
|
||||
"Add two numbers together and return the result."
|
||||
)
|
||||
add_numbers.schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {"type": "number", "description": "First number"},
|
||||
"b": {"type": "number", "description": "Second number"},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
}
|
||||
|
||||
# Use the tool
|
||||
result, _ = await add_numbers(a=5, b=7)
|
||||
print(f"\nCustom tool result: {result}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run both examples
|
||||
loop = asyncio.get_event_loop()
|
||||
loop.run_until_complete(main())
|
||||
loop.run_until_complete(create_custom_tool())
|
||||
Loading…
Reference in new issue