google-adk 0.5.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

google/adk/tools/mcp_tool/mcp_toolset.py

@@ -12,255 +12,160 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from contextlib import AsyncExitStack
+import logging
 import sys
-from types import TracebackType
-from typing import List, Optional, TextIO, Tuple, Type
-
-from .mcp_session_manager import MCPSessionManager, SseServerParams, retry_on_closed_resource
+from typing import List
+from typing import Optional
+from typing import TextIO
+from typing import Union
+
+from ...agents.readonly_context import ReadonlyContext
+from ..base_tool import BaseTool
+from ..base_toolset import BaseToolset
+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
 
 # Attempt to import MCP Tool from the MCP library, and hints user to upgrade
 # their Python version to 3.10 if it fails.
 try:
-  from mcp import ClientSession, StdioServerParameters
+  from mcp import StdioServerParameters
   from mcp.types import ListToolsResult
 except ImportError as e:
   import sys
 
   if sys.version_info < (3, 10):
     raise ImportError(
-        'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
-        ' version.'
+        "MCP Tool requires Python 3.10 or above. Please upgrade your Python"
+        " version."
     ) from e
   else:
     raise e
 
 from .mcp_tool import MCPTool
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
-class MCPToolset:
+class MCPToolset(BaseToolset):
   """Connects to a MCP Server, and retrieves MCP Tools into ADK Tools.
 
+  This toolset manages the connection to an MCP server and provides tools
+  that can be used by an agent. It properly implements the BaseToolset
+  interface for easy integration with the agent framework.
+
   Usage:
-  Example 1: (using from_server helper):
-  ```
-  async def load_tools():
-    return await MCPToolset.from_server(
+  ```python
+  toolset = MCPToolset(
       connection_params=StdioServerParameters(
           command='npx',
           args=["-y", "@modelcontextprotocol/server-filesystem"],
-          )
-    )
+      ),
+      tool_filter=['read_file', 'list_directory']  # Optional: filter specific tools
+  )
 
-  # Use the tools in an LLM agent
-  tools, exit_stack = await load_tools()
+  # Use in an agent
   agent = LlmAgent(
-      tools=tools
+      model='gemini-2.0-flash',
+      name='enterprise_assistant',
+      instruction='Help user accessing their file systems',
+      tools=[toolset],
   )
-  ...
-  await exit_stack.aclose()
-  ```
-
-  Example 2: (using `async with`):
-
-  ```
-  async def load_tools():
-    async with MCPToolset(
-      connection_params=SseServerParams(url="http://0.0.0.0:8090/sse")
-    ) as toolset:
-      tools = await toolset.load_tools()
-
-      agent = LlmAgent(
-          ...
-          tools=tools
-      )
-  ```
-
-  Example 3: (provide AsyncExitStack):
-  ```
-  async def load_tools():
-    async_exit_stack = AsyncExitStack()
-    toolset = MCPToolset(
-      connection_params=StdioServerParameters(...),
-    )
-    async_exit_stack.enter_async_context(toolset)
-    tools = await toolset.load_tools()
-    agent = LlmAgent(
-        ...
-        tools=tools
-    )
-    ...
-    await async_exit_stack.aclose()
 
+  # Cleanup is handled automatically by the agent framework
+  # But you can also manually close if needed:
+  # await toolset.close()
   ```
-
-  Attributes:
-    connection_params: The connection parameters to the MCP server. Can be
-      either `StdioServerParameters` or `SseServerParams`.
-    exit_stack: The async exit stack to manage the connection to the MCP server.
-    session: The MCP session being initialized with the connection.
   """
 
   def __init__(
       self,
       *,
       connection_params: StdioServerParameters | SseServerParams,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
       errlog: TextIO = sys.stderr,
-      exit_stack=AsyncExitStack(),
   ):
     """Initializes the MCPToolset.
 
-    Usage:
-    Example 1: (using from_server helper):
-    ```
-    async def load_tools():
-      return await MCPToolset.from_server(
-        connection_params=StdioServerParameters(
-            command='npx',
-            args=["-y", "@modelcontextprotocol/server-filesystem"],
-            )
-      )
-
-    # Use the tools in an LLM agent
-    tools, exit_stack = await load_tools()
-    agent = LlmAgent(
-        tools=tools
-    )
-    ...
-    await exit_stack.aclose()
-    ```
-
-    Example 2: (using `async with`):
-
-    ```
-    async def load_tools():
-      async with MCPToolset(
-        connection_params=SseServerParams(url="http://0.0.0.0:8090/sse")
-      ) as toolset:
-        tools = await toolset.load_tools()
-
-        agent = LlmAgent(
-            ...
-            tools=tools
-        )
-    ```
-
-    Example 3: (provide AsyncExitStack):
-    ```
-    async def load_tools():
-      async_exit_stack = AsyncExitStack()
-      toolset = MCPToolset(
-        connection_params=StdioServerParameters(...),
-      )
-      async_exit_stack.enter_async_context(toolset)
-      tools = await toolset.load_tools()
-      agent = LlmAgent(
-          ...
-          tools=tools
-      )
-      ...
-      await async_exit_stack.aclose()
-
-    ```
-
     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.
+        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:
+            - A list of tool names to include
+            - A ToolPredicate function for custom filtering logic
+        errlog: TextIO stream for error logging.
     """
+    super().__init__(tool_filter=tool_filter)
+
     if not connection_params:
-      raise ValueError('Missing connection params in MCPToolset.')
-    self.connection_params = connection_params
-    self.errlog = errlog
-    self.exit_stack = exit_stack
+      raise ValueError("Missing connection params in MCPToolset.")
 
-    self.session_manager = MCPSessionManager(
-        connection_params=self.connection_params,
-        exit_stack=self.exit_stack,
-        errlog=self.errlog,
-    )
+    self._connection_params = connection_params
+    self._errlog = errlog
 
-  @classmethod
-  async def from_server(
-      cls,
-      *,
-      connection_params: StdioServerParameters | SseServerParams,
-      async_exit_stack: Optional[AsyncExitStack] = None,
-      errlog: TextIO = sys.stderr,
-  ) -> Tuple[List[MCPTool], AsyncExitStack]:
-    """Retrieve all tools from the MCP connection.
+    # Create the session manager that will handle the MCP connection
+    self._mcp_session_manager = MCPSessionManager(
+        connection_params=self._connection_params,
+        errlog=self._errlog,
+    )
+    self._session = None
 
-    Usage:
-    ```
-    async def load_tools():
-      tools, exit_stack = await MCPToolset.from_server(
-        connection_params=StdioServerParameters(
-            command='npx',
-            args=["-y", "@modelcontextprotocol/server-filesystem"],
-        )
-      )
-    ```
+  @retry_on_closed_resource("_reinitialize_session")
+  async def get_tools(
+      self,
+      readonly_context: Optional[ReadonlyContext] = None,
+  ) -> List[BaseTool]:
+    """Return all tools in the toolset based on the provided context.
 
     Args:
-      connection_params: The connection parameters to the MCP server.
-      async_exit_stack: The async exit stack to use. If not provided, a new
-        AsyncExitStack will be created.
+        readonly_context: Context used to filter tools available to the agent.
+            If None, all tools in the toolset are returned.
 
     Returns:
-      A tuple of the list of MCPTools and the AsyncExitStack.
-      - tools: The list of MCPTools.
-      - async_exit_stack: The AsyncExitStack used to manage the connection to
-        the MCP server. Use `await async_exit_stack.aclose()` to close the
-        connection when server shuts down.
+        List[BaseTool]: A list of tools available under the specified context.
     """
-    async_exit_stack = async_exit_stack or AsyncExitStack()
-    toolset = cls(
-        connection_params=connection_params,
-        exit_stack=async_exit_stack,
-        errlog=errlog,
-    )
+    # Get session from session manager
+    if not self._session:
+      self._session = await self._mcp_session_manager.create_session()
+
+    # Fetch available tools from the MCP server
+    tools_response: ListToolsResult = await self._session.list_tools()
+
+    # Apply filtering based on context and tool_filter
+    tools = []
+    for tool in tools_response.tools:
+      mcp_tool = MCPTool(
+          mcp_tool=tool,
+          mcp_session_manager=self._mcp_session_manager,
+      )
 
-    await async_exit_stack.enter_async_context(toolset)
-    tools = await toolset.load_tools()
-    return (tools, async_exit_stack)
+      if self._is_tool_selected(mcp_tool, readonly_context):
+        tools.append(mcp_tool)
+    return tools
 
-  async def _initialize(self) -> ClientSession:
-    """Connects to the MCP Server and initializes the ClientSession."""
-    self.session = await self.session_manager.create_session()
-    return self.session
+  async def _reinitialize_session(self):
+    """Reinitializes the session when connection is lost."""
+    # Close the old session and clear cache
+    await self._mcp_session_manager.close()
+    self._session = await self._mcp_session_manager.create_session()
 
-  async def _exit(self):
-    """Closes the connection to MCP Server."""
-    await self.exit_stack.aclose()
+    # Tools will be reloaded on next get_tools call
 
-  @retry_on_closed_resource('_initialize')
-  async def load_tools(self) -> List[MCPTool]:
-    """Loads all tools from the MCP Server.
+  async def close(self) -> None:
+    """Performs cleanup and releases resources held by the toolset.
 
-    Returns:
-      A list of MCPTools imported from the MCP Server.
+    This method closes the MCP session and cleans up all associated resources.
+    It's designed to be safe to call multiple times and handles cleanup errors
+    gracefully to avoid blocking application shutdown.
     """
-    tools_response: ListToolsResult = await self.session.list_tools()
-    return [
-        MCPTool(
-            mcp_tool=tool,
-            mcp_session=self.session,
-            mcp_session_manager=self.session_manager,
-        )
-        for tool in tools_response.tools
-    ]
-
-  async def __aenter__(self):
     try:
-      await self._initialize()
-      return self
+      await self._mcp_session_manager.close()
     except Exception as e:
-      raise e
-
-  async def __aexit__(
-      self,
-      exc_type: Optional[Type[BaseException]],
-      exc: Optional[BaseException],
-      tb: Optional[TracebackType],
-  ) -> None:
-    await self._exit()
+      # Log the error but don't re-raise to avoid blocking shutdown
+      print(f"Warning: Error during MCPToolset cleanup: {e}", file=self._errlog)
+    finally:
+      # Clear cached tools
+      self._tools_cache = None
+      self._tools_loaded = False

google/adk/tools/openapi_tool/auth/credential_exchangers/base_credential_exchanger.py

@@ -15,9 +15,7 @@
 import abc
 from typing import Optional
 
-from .....auth.auth_credential import (
-    AuthCredential,
-)
+from .....auth.auth_credential import AuthCredential
 from .....auth.auth_schemes import AuthScheme
 
 

google/adk/tools/openapi_tool/auth/credential_exchangers/service_account_exchanger.py

@@ -21,14 +21,13 @@ from google.auth.transport.requests import Request
 from google.oauth2 import service_account
 import google.oauth2.credentials
 
-from .....auth.auth_credential import (
-    AuthCredential,
-    AuthCredentialTypes,
-    HttpAuth,
-    HttpCredentials,
-)
+from .....auth.auth_credential import AuthCredential
+from .....auth.auth_credential import AuthCredentialTypes
+from .....auth.auth_credential import HttpAuth
+from .....auth.auth_credential import HttpCredentials
 from .....auth.auth_schemes import AuthScheme
-from .base_credential_exchanger import AuthCredentialMissingError, BaseAuthCredentialExchanger
+from .base_credential_exchanger import AuthCredentialMissingError
+from .base_credential_exchanger import BaseAuthCredentialExchanger
 
 
 class ServiceAccountCredentialExchanger(BaseAuthCredentialExchanger):

google/adk/tools/openapi_tool/common/common.py

@@ -14,7 +14,11 @@
 
 import keyword
 import re
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
 
 from fastapi.openapi.models import Response
 from fastapi.openapi.models import Schema

google/adk/tools/openapi_tool/openapi_spec_parser/__init__.py

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .openapi_spec_parser import OpenApiSpecParser, OperationEndpoint, ParsedOperation
+from .openapi_spec_parser import OpenApiSpecParser
+from .openapi_spec_parser import OperationEndpoint
+from .openapi_spec_parser import ParsedOperation
 from .openapi_toolset import OpenAPIToolset
 from .operation_parser import OperationParser
-from .rest_api_tool import AuthPreparationState, RestApiTool, snake_to_lower_camel, to_gemini_schema
+from .rest_api_tool import AuthPreparationState
+from .rest_api_tool import RestApiTool
+from .rest_api_tool import snake_to_lower_camel
+from .rest_api_tool import to_gemini_schema
 from .tool_auth_handler import ToolAuthHandler
 
 __all__ = [

google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py

@@ -20,18 +20,23 @@ from typing import Final
 from typing import List
 from typing import Literal
 from typing import Optional
+from typing import Union
 
+from typing_extensions import override
 import yaml
 
+from ....agents.readonly_context import ReadonlyContext
 from ....auth.auth_credential import AuthCredential
 from ....auth.auth_schemes import AuthScheme
+from ...base_toolset import BaseToolset
+from ...base_toolset import ToolPredicate
 from .openapi_spec_parser import OpenApiSpecParser
 from .rest_api_tool import RestApiTool
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 
-class OpenAPIToolset:
+class OpenAPIToolset(BaseToolset):
   """Class for parsing OpenAPI spec into a list of RestApiTool.
 
   Usage:
@@ -61,6 +66,7 @@ class OpenAPIToolset:
       spec_str_type: Literal["json", "yaml"] = "json",
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
     """Initializes the OpenAPIToolset.
 
@@ -94,10 +100,13 @@ class OpenAPIToolset:
       auth_credential: The auth credential to use for all tools. Use
         AuthCredential or use helpers in
         `google.adk.tools.openapi_tool.auth.auth_helpers`
+      tool_filter: The filter used to filter the tools in the toolset. It can be
+        either a tool predicate or a list of tool names of the tools to expose.
     """
+    super().__init__(tool_filter=tool_filter)
     if not spec_dict:
       spec_dict = self._load_spec(spec_str, spec_str_type)
-    self.tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
+    self._tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
     if auth_scheme or auth_credential:
       self._configure_auth_all(auth_scheme, auth_credential)
 
@@ -106,19 +115,26 @@ class OpenAPIToolset:
   ):
     """Configure auth scheme and credential for all tools."""
 
-    for tool in self.tools:
+    for tool in self._tools:
       if auth_scheme:
         tool.configure_auth_scheme(auth_scheme)
       if auth_credential:
         tool.configure_auth_credential(auth_credential)
 
-  def get_tools(self) -> List[RestApiTool]:
+  @override
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> List[RestApiTool]:
     """Get all tools in the toolset."""
-    return self.tools
+    return [
+        tool
+        for tool in self._tools
+        if self._is_tool_selected(tool, readonly_context)
+    ]
 
   def get_tool(self, tool_name: str) -> Optional[RestApiTool]:
     """Get a tool by name."""
-    matching_tool = filter(lambda t: t.name == tool_name, self.tools)
+    matching_tool = filter(lambda t: t.name == tool_name, self._tools)
     return next(matching_tool, None)
 
   def _load_spec(
@@ -142,3 +158,7 @@ class OpenAPIToolset:
       logger.info("Parsed tool: %s", tool.name)
       tools.append(tool)
     return tools
+
+  @override
+  async def close(self):
+    pass