mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
80b0920e58
commit
ad3d8171bf
@@ -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 {},
|
||||
|
||||
Reference in New Issue
Block a user