Python: Introducing Local MCP Servers (#389)

* mcp parts

* mcp parts 2

* removed structured output in favor of handling in chatresponse, mcp as AITool and running samples

* updated naming

* fixed test
This commit is contained in:
Eduard van Valkenburg
2025-08-13 11:48:22 +02:00
committed by GitHub
Unverified
parent 80b0920e58
commit ad3d8171bf
20 changed files with 1970 additions and 298 deletions
+50 -11
View File
@@ -2,14 +2,16 @@
import sys
from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
from contextlib import AbstractAsyncContextManager
from contextlib import AbstractAsyncContextManager, AsyncExitStack
from enum import Enum
from itertools import chain
from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
from uuid import uuid4
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, PrivateAttr
from ._clients import ChatClient
from ._mcp import McpTool
from ._pydantic import AFBaseModel
from ._tools import AITool
from ._types import (
@@ -315,6 +317,8 @@ class ChatClientAgent(AgentBase):
chat_client: ChatClient
instructions: str | None = None
chat_options: ChatOptions
_local_mcp_tools: list[McpTool] = PrivateAttr(default_factory=list) # type: ignore[reportUnknownVariableType]
_async_exit_stack: AsyncExitStack = PrivateAttr(default_factory=AsyncExitStack)
def __init__(
self,
@@ -383,6 +387,11 @@ class ChatClientAgent(AgentBase):
"""
kwargs.update(additional_properties or {})
# We ignore the MCP Servers here and store them separately,
# we add their functions to the tools list at runtime
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, McpTool)]
final_tools = [tool for tool in normalized_tools if not isinstance(tool, McpTool)]
args: dict[str, Any] = {
"chat_client": chat_client,
"chat_options": ChatOptions(
@@ -398,7 +407,7 @@ class ChatClientAgent(AgentBase):
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=tools, # type: ignore
tools=final_tools, # type: ignore[reportArgumentType]
top_p=top_p,
user=user,
additional_properties=kwargs,
@@ -415,23 +424,27 @@ class ChatClientAgent(AgentBase):
super().__init__(**args)
self._update_agent_name()
self._local_mcp_tools = local_mcp_tools # type: ignore[assignment]
async def __aenter__(self) -> "Self":
"""Async context manager entry.
If the chat_client supports async context management, enter its context.
If either the chat_client or the local_mcp_tools are context managers,
they will be entered into the async exit stack to ensure proper cleanup.
This list might be extended in the future.
"""
if isinstance(self.chat_client, AbstractAsyncContextManager):
await self.chat_client.__aenter__() # type: ignore[reportUnknownMemberType]
for context_manager in chain([self.chat_client], self._local_mcp_tools):
if isinstance(context_manager, AbstractAsyncContextManager):
await self._async_exit_stack.enter_async_context(context_manager)
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Async context manager exit.
If the chat_client supports async context management, exit its context.
Close the async exit stack to ensure all context managers are exited properly.
"""
if isinstance(self.chat_client, AbstractAsyncContextManager):
await self.chat_client.__aexit__(exc_type, exc_val, exc_tb) # type: ignore[reportUnknownMemberType]
await self._async_exit_stack.aclose()
def _update_agent_name(self) -> None:
"""Update the agent name in a chat client.
@@ -506,6 +519,19 @@ class ChatClientAgent(AgentBase):
thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
agent_name = self._get_agent_name()
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[AITool | dict[str, Any] | Callable[..., Any]] = []
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, McpTool):
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool) # type: ignore
for mcp_server in self._local_mcp_tools:
final_tools.extend(mcp_server.functions)
response = await self.chat_client.get_response(
messages=thread_messages,
chat_options=self.chat_options
@@ -523,7 +549,7 @@ class ChatClientAgent(AgentBase):
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=tools, # type: ignore
tools=final_tools, # type: ignore[reportArgumentType]
top_p=top_p,
user=user,
additional_properties=additional_properties or {},
@@ -617,6 +643,19 @@ class ChatClientAgent(AgentBase):
agent_name = self._get_agent_name()
response_updates: list[ChatResponseUpdate] = []
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[AITool | MutableMapping[str, Any] | Callable[..., Any]] = []
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, McpTool):
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool)
for mcp_server in self._local_mcp_tools:
final_tools.extend(mcp_server.functions)
async for update in self.chat_client.get_streaming_response(
messages=thread_messages,
chat_options=self.chat_options
@@ -634,7 +673,7 @@ class ChatClientAgent(AgentBase):
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=tools, # type: ignore
tools=final_tools, # type: ignore[reportArgumentType]
top_p=top_p,
user=user,
additional_properties=additional_properties or {},