adapt google api toolset and api hub toolset to new toolset interface

PiperOrigin-RevId: 757946541
This commit is contained in:
Xiang (Sean) Zhou
2025-05-12 15:57:56 -07:00
committed by Copybara-Service
parent 27b229719e
commit 6a04ff84ba
7 changed files with 168 additions and 122 deletions

View File

@@ -14,7 +14,7 @@
from contextlib import AsyncExitStack
import sys
from typing import List
from typing import List, Union
from typing import Optional
from typing import TextIO
@@ -68,7 +68,7 @@ class MCPToolset(BaseToolset):
*,
connection_params: StdioServerParameters | SseServerParams,
errlog: TextIO = sys.stderr,
tool_predicate: Optional[ToolPredicate] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
):
"""Initializes the MCPToolset.
@@ -90,7 +90,7 @@ class MCPToolset(BaseToolset):
errlog=self.errlog,
)
self.session = None
self.tool_predicate = tool_predicate
self.tool_filter = tool_filter
async def _initialize(self) -> ClientSession:
"""Connects to the MCP Server and initializes the ClientSession."""
@@ -106,7 +106,7 @@ class MCPToolset(BaseToolset):
@override
async def get_tools(
self,
readony_context: ReadonlyContext = None,
readonly_context: Optional[ReadonlyContext] = None,
) -> List[MCPTool]:
"""Loads all tools from the MCP Server.
@@ -123,6 +123,5 @@ class MCPToolset(BaseToolset):
mcp_session_manager=self.session_manager,
)
for tool in tools_response.tools
if self.tool_predicate is None
or self.tool_predicate(tool, readony_context)
if self.tool_filter is None or self.tool_filter(tool, readonly_context)
]