mirror of
https://github.com/EvolutionAPI/adk-python.git
synced 2025-12-25 14:37:45 -06:00
feat: google/adk-python#479 support for streamable http MCP servers for MCPToolset
Copybara import of the project: -- c5b9d49d7b6d858ff0a93bd690e6d653b7c32221 by Omar BENHAMID <omar.benhamid@smart-gts.com>: feat: google/adk-python#479 support for streamable http MCP servers for MCPToolset -- 9431bc19e6538c1b814aba0b24ff564acf046075 by Omar BENHAMID <omar.benhamid@smart-gts.com>: feat: google/adk-python#479 streamable http added to right package -- 8b4aabed45a6f0dc828beb61f12985dc7b14f3d0 by Omar BENHAMID <omar.benhamid@smart-gts.com>: feat: google/adk-python#479 streamable http : review feedbacks + sample agent COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/650 from omarbenhamid:feature/mcp-streamable-http 625f028784c216401d45cb1b5d4d998535ebcb00 PiperOrigin-RevId: 764419586
This commit is contained in:
committed by
Copybara-Service
parent
c7ce987676
commit
d232e6216d
@@ -17,10 +17,10 @@ from contextlib import AsyncExitStack
|
||||
import functools
|
||||
import logging
|
||||
import sys
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
from typing import Optional
|
||||
from typing import TextIO
|
||||
|
||||
import anyio
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -29,6 +29,7 @@ try:
|
||||
from mcp import StdioServerParameters
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
except ImportError as e:
|
||||
import sys
|
||||
|
||||
@@ -56,6 +57,19 @@ class SseServerParams(BaseModel):
|
||||
sse_read_timeout: float = 60 * 5
|
||||
|
||||
|
||||
class StreamableHTTPServerParams(BaseModel):
|
||||
"""Parameters for the MCP SSE connection.
|
||||
|
||||
See MCP SSE Client documentation for more details.
|
||||
https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py
|
||||
"""
|
||||
url: str
|
||||
headers: dict[str, Any] | None = None
|
||||
timeout: float = 5
|
||||
sse_read_timeout: float = 60 * 5
|
||||
terminate_on_close: bool = True
|
||||
|
||||
|
||||
def retry_on_closed_resource(async_reinit_func_name: str):
|
||||
"""Decorator to automatically reinitialize session and retry action.
|
||||
|
||||
@@ -123,13 +137,13 @@ class MCPSessionManager:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connection_params: StdioServerParameters | SseServerParams,
|
||||
connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
|
||||
errlog: TextIO = sys.stderr,
|
||||
):
|
||||
"""Initializes the MCP session manager.
|
||||
|
||||
Args:
|
||||
connection_params: Parameters for the MCP connection (Stdio or SSE).
|
||||
connection_params: Parameters for the MCP connection (Stdio, SSE or Streamable HTTP).
|
||||
errlog: (Optional) TextIO stream for error logging. Use only for
|
||||
initializing a local stdio MCP session.
|
||||
"""
|
||||
@@ -163,6 +177,14 @@ class MCPSessionManager:
|
||||
timeout=self._connection_params.timeout,
|
||||
sse_read_timeout=self._connection_params.sse_read_timeout,
|
||||
)
|
||||
elif isinstance(self._connection_params, StreamableHTTPServerParams):
|
||||
client = streamablehttp_client(
|
||||
url=self._connection_params.url,
|
||||
headers=self._connection_params.headers,
|
||||
timeout=timedelta(seconds=self._connection_params.timeout),
|
||||
sse_read_timeout=timedelta(seconds=self._connection_params.sse_read_timeout),
|
||||
terminate_on_close=self._connection_params.terminate_on_close,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
'Unable to initialize connection. Connection should be'
|
||||
@@ -171,8 +193,10 @@ class MCPSessionManager:
|
||||
)
|
||||
|
||||
transports = await self._exit_stack.enter_async_context(client)
|
||||
# The streamable http client returns a GetSessionCallback in addition to the read/write MemoryObjectStreams
|
||||
# needed to build the ClientSession, we limit then to the two first values to be compatible with all clients.
|
||||
session = await self._exit_stack.enter_async_context(
|
||||
ClientSession(*transports)
|
||||
ClientSession(*transports[:2])
|
||||
)
|
||||
await session.initialize()
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@ from ..base_toolset import ToolPredicate
|
||||
from .mcp_session_manager import MCPSessionManager
|
||||
from .mcp_session_manager import retry_on_closed_resource
|
||||
from .mcp_session_manager import SseServerParams
|
||||
from .mcp_session_manager import StreamableHTTPServerParams
|
||||
|
||||
# Attempt to import MCP Tool from the MCP library, and hints user to upgrade
|
||||
# their Python version to 3.10 if it fails.
|
||||
@@ -82,17 +83,18 @@ class MCPToolset(BaseToolset):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
connection_params: StdioServerParameters | SseServerParams,
|
||||
connection_params: StdioServerParameters | SseServerParams | StreamableHTTPServerParams,
|
||||
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
|
||||
errlog: TextIO = sys.stderr,
|
||||
):
|
||||
"""Initializes the MCPToolset.
|
||||
|
||||
Args:
|
||||
connection_params: The connection parameters to the MCP server. Can be:
|
||||
`StdioServerParameters` for using local mcp server (e.g. using `npx` or
|
||||
`python3`); or `SseServerParams` for a local/remote SSE server.
|
||||
tool_filter: Optional filter to select specific tools. Can be either:
|
||||
connection_params: The connection parameters to the MCP server. Can be:
|
||||
`StdioServerParameters` for using local mcp server (e.g. using `npx` or
|
||||
`python3`); or `SseServerParams` for a local/remote SSE server; or
|
||||
`StreamableHTTPServerParams` for local/remote Streamable http server.
|
||||
tool_filter: Optional filter to select specific tools. Can be either:
|
||||
- A list of tool names to include
|
||||
- A ToolPredicate function for custom filtering logic
|
||||
errlog: TextIO stream for error logging.
|
||||
@@ -110,6 +112,7 @@ class MCPToolset(BaseToolset):
|
||||
connection_params=self._connection_params,
|
||||
errlog=self._errlog,
|
||||
)
|
||||
|
||||
self._session = None
|
||||
|
||||
@retry_on_closed_resource("_reinitialize_session")
|
||||
|
||||
Reference in New Issue
Block a user