google-adk 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

google/adk/tools/bigquery/query_tool.py

@@ -0,0 +1,76 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.oauth2.credentials import Credentials
+
+from ...tools.bigquery import client
+
+MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50
+
+
+def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
+  """Run a BigQuery SQL query in the project and return the result.
+
+  Args:
+      project_id (str): The GCP project id in which the query should be
+        executed.
+      query (str): The BigQuery SQL query to be executed.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary representing the result of the query.
+            If the result contains the key "result_is_likely_truncated" with
+            value True, it means that there may be additional rows matching the
+            query not returned in the result.
+
+  Examples:
+      >>> execute_sql("bigframes-dev",
+      ... "SELECT island, COUNT(*) AS population "
+      ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
+      {
+        "rows": [
+            {
+                "island": "Dream",
+                "population": 124
+            },
+            {
+                "island": "Biscoe",
+                "population": 168
+            },
+            {
+                "island": "Torgersen",
+                "population": 52
+            }
+        ]
+      }
+  """
+
+  try:
+    bq_client = client.get_bigquery_client(credentials=credentials)
+    row_iterator = bq_client.query_and_wait(
+        query, project=project_id, max_results=MAX_DOWNLOADED_QUERY_RESULT_ROWS
+    )
+    rows = [{key: val for key, val in row.items()} for row in row_iterator]
+    result = {"rows": rows}
+    if (
+        MAX_DOWNLOADED_QUERY_RESULT_ROWS is not None
+        and len(rows) == MAX_DOWNLOADED_QUERY_RESULT_ROWS
+    ):
+      result["result_is_likely_truncated"] = True
+    return result
+  except Exception as ex:
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }

google/adk/tools/function_parameter_parse_util.py

@@ -289,6 +289,13 @@ def _parse_schema_from_parameter(
       )
     _raise_if_schema_unsupported(variant, schema)
     return schema
+  if param.annotation is None:
+    # https://swagger.io/docs/specification/v3_0/data-models/data-types/#null
+    # null is not a valid type in schema, use object instead.
+    schema.type = types.Type.OBJECT
+    schema.nullable = True
+    _raise_if_schema_unsupported(variant, schema)
+    return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'
       ' automatic function calling. Automatic function calling works best with'

google/adk/tools/function_tool.py

@@ -33,8 +33,31 @@ class FunctionTool(BaseTool):
   """
 
   def __init__(self, func: Callable[..., Any]):
-    super().__init__(name=func.__name__, description=func.__doc__)
+    """Extract metadata from a callable object."""
+    name = ''
+    doc = ''
+    # Handle different types of callables
+    if hasattr(func, '__name__'):
+      # Regular functions, unbound methods, etc.
+      name = func.__name__
+    elif hasattr(func, '__class__'):
+      # Callable objects, bound methods, etc.
+      name = func.__class__.__name__
+
+    # Get documentation (prioritize direct __doc__ if available)
+    if hasattr(func, '__doc__') and func.__doc__:
+      doc = func.__doc__
+    elif (
+        hasattr(func, '__call__')
+        and hasattr(func.__call__, '__doc__')
+        and func.__call__.__doc__
+    ):
+      # For callable objects, try to get docstring from __call__ method
+      doc = func.__call__.__doc__
+
+    super().__init__(name=name, description=doc)
     self.func = func
+    self._ignore_params = ['tool_context', 'input_stream']
 
   @override
   def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
@@ -43,7 +66,7 @@ class FunctionTool(BaseTool):
             func=self.func,
             # The model doesn't understand the function context.
             # input_stream is for streaming tool
-            ignore_params=['tool_context', 'input_stream'],
+            ignore_params=self._ignore_params,
             variant=self._api_variant,
         )
     )
@@ -76,7 +99,14 @@ class FunctionTool(BaseTool):
 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters."""
       return {'error': error_str}
 
-    if inspect.iscoroutinefunction(self.func):
+    # Functions are callable objects, but not all callable objects are functions
+    # checking coroutine function is not enough. We also need to check whether
+    # Callable's __call__ function is a coroutine funciton
+    if (
+        inspect.iscoroutinefunction(self.func)
+        or hasattr(self.func, '__call__')
+        and inspect.iscoroutinefunction(self.func.__call__)
+    ):
       return await self.func(**args_to_call) or {}
     else:
       return self.func(**args_to_call) or {}

google/adk/tools/get_user_choice_tool.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import Optional
+
 from .long_running_tool import LongRunningFunctionTool
 from .tool_context import ToolContext
 

google/adk/tools/google_api_tool/__init__.py

@@ -11,18 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__all__ = [
-    'BigQueryToolset',
-    'CalendarToolset',
-    'GmailToolset',
-    'YoutubeToolset',
-    'SlidesToolset',
-    'SheetsToolset',
-    'DocsToolset',
-    'GoogleApiToolset',
-    'GoogleApiTool',
-]
 
+"""Auto-generated tools and toolsets for Google APIs.
+
+These tools and toolsets are auto-generated based on the API specifications
+provided by the Google API Discovery API.
+"""
 
 from .google_api_tool import GoogleApiTool
 from .google_api_toolset import GoogleApiToolset
@@ -33,3 +27,15 @@ from .google_api_toolsets import GmailToolset
 from .google_api_toolsets import SheetsToolset
 from .google_api_toolsets import SlidesToolset
 from .google_api_toolsets import YoutubeToolset
+
+__all__ = [
+    'BigQueryToolset',
+    'CalendarToolset',
+    'GmailToolset',
+    'YoutubeToolset',
+    'SlidesToolset',
+    'SheetsToolset',
+    'DocsToolset',
+    'GoogleApiToolset',
+    'GoogleApiTool',
+]

google/adk/tools/google_api_tool/google_api_tool.py

@@ -19,10 +19,10 @@ from typing import Optional
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
+from .. import BaseTool
 from ...auth import AuthCredential
 from ...auth import AuthCredentialTypes
 from ...auth import OAuth2Auth
-from .. import BaseTool
 from ..openapi_tool import RestApiTool
 from ..tool_context import ToolContext
 

google/adk/tools/google_api_tool/google_api_toolset.py

@@ -56,20 +56,6 @@ class GoogleApiToolset(BaseToolset):
     self._openapi_toolset = self._load_toolset_with_oidc_auth()
     self.tool_filter = tool_filter
 
-  def _is_tool_selected(
-      self, tool: GoogleApiTool, readonly_context: ReadonlyContext
-  ) -> bool:
-    if not self.tool_filter:
-      return True
-
-    if isinstance(self.tool_filter, ToolPredicate):
-      return self.tool_filter(tool, readonly_context)
-
-    if isinstance(self.tool_filter, list):
-      return tool.name in self.tool_filter
-
-    return False
-
   @override
   async def get_tools(
       self, readonly_context: Optional[ReadonlyContext] = None

google/adk/tools/google_api_tool/google_api_toolsets.py

@@ -18,14 +18,14 @@ from typing import List
 from typing import Optional
 from typing import Union
 
-from google.adk.tools.base_toolset import ToolPredicate
-
+from ..base_toolset import ToolPredicate
 from .google_api_toolset import GoogleApiToolset
 
 logger = logging.getLogger("google_adk." + __name__)
 
 
 class BigQueryToolset(GoogleApiToolset):
+  """Auto-generated Bigquery toolset based on Google BigQuery API v2 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -37,6 +37,7 @@ class BigQueryToolset(GoogleApiToolset):
 
 
 class CalendarToolset(GoogleApiToolset):
+  """Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -48,6 +49,7 @@ class CalendarToolset(GoogleApiToolset):
 
 
 class GmailToolset(GoogleApiToolset):
+  """Auto-generated Gmail toolset based on Google Gmail API v1 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -59,6 +61,7 @@ class GmailToolset(GoogleApiToolset):
 
 
 class YoutubeToolset(GoogleApiToolset):
+  """Auto-generated Youtube toolset based on Youtube API v3 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -70,6 +73,7 @@ class YoutubeToolset(GoogleApiToolset):
 
 
 class SlidesToolset(GoogleApiToolset):
+  """Auto-generated Slides toolset based on Google Slides API v1 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -81,6 +85,7 @@ class SlidesToolset(GoogleApiToolset):
 
 
 class SheetsToolset(GoogleApiToolset):
+  """Auto-generated Sheets toolset based on Google Sheets API v4 spec exposed by Google API discovery API"""
 
   def __init__(
       self,
@@ -92,6 +97,7 @@ class SheetsToolset(GoogleApiToolset):
 
 
 class DocsToolset(GoogleApiToolset):
+  """Auto-generated Docs toolset based on Google Docs API v1 spec exposed by Google API discovery API"""
 
   def __init__(
       self,

google/adk/tools/google_search_tool.py

@@ -46,7 +46,7 @@ class GoogleSearchTool(BaseTool):
   ) -> None:
     llm_request.config = llm_request.config or types.GenerateContentConfig()
     llm_request.config.tools = llm_request.config.tools or []
-    if llm_request.model and llm_request.model.startswith('gemini-1'):
+    if llm_request.model and 'gemini-1' in llm_request.model:
       if llm_request.config.tools:
         print(llm_request.config.tools)
         raise ValueError(
@@ -55,7 +55,7 @@ class GoogleSearchTool(BaseTool):
       llm_request.config.tools.append(
           types.Tool(google_search_retrieval=types.GoogleSearchRetrieval())
       )
-    elif llm_request.model and llm_request.model.startswith('gemini-2'):
+    elif llm_request.model and 'gemini-2' in llm_request.model:
       llm_request.config.tools.append(
           types.Tool(google_search=types.GoogleSearch())
       )

google/adk/tools/mcp_tool/conversion_utils.py

@@ -12,9 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict
-from google.genai.types import Schema, Type
+from typing import Any
+from typing import Dict
+
+from google.genai.types import Schema
+from google.genai.types import Type
 import mcp.types as mcp_types
+
 from ..base_tool import BaseTool
 
 

google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
-from contextlib import asynccontextmanager
+
 from contextlib import AsyncExitStack
 import functools
 import logging
@@ -71,29 +70,27 @@ def retry_on_closed_resource(async_reinit_func_name: str):
 
   Usage:
   class MCPTool:
-    ...
-    async def create_session(self):
-      self.session = ...
+      ...
+      async def create_session(self):
+          self.session = ...
 
-    @retry_on_closed_resource('create_session')
-    async def use_session(self):
-      await self.session.call_tool()
+      @retry_on_closed_resource('create_session')
+      async def use_session(self):
+          await self.session.call_tool()
 
   Args:
-    async_reinit_func_name: The name of the async function to recreate session.
+      async_reinit_func_name: The name of the async function to recreate session.
 
   Returns:
-    The decorated function.
+      The decorated function.
   """
 
   def decorator(func):
-    @functools.wraps(
-        func
-    )  # Preserves original function metadata (name, docstring)
+    @functools.wraps(func)  # Preserves original function metadata
     async def wrapper(self, *args, **kwargs):
       try:
         return await func(self, *args, **kwargs)
-      except anyio.ClosedResourceError:
+      except anyio.ClosedResourceError as close_err:
         try:
           if hasattr(self, async_reinit_func_name) and callable(
               getattr(self, async_reinit_func_name)
@@ -105,7 +102,7 @@ def retry_on_closed_resource(async_reinit_func_name: str):
                 f'Function {async_reinit_func_name} does not exist in decorated'
                 ' class. Please check the function name in'
                 ' retry_on_closed_resource decorator.'
-            )
+            ) from close_err
         except Exception as reinit_err:
           raise RuntimeError(
               f'Error reinitializing: {reinit_err}'
@@ -117,45 +114,6 @@ def retry_on_closed_resource(async_reinit_func_name: str):
   return decorator
 
 
-@asynccontextmanager
-async def tracked_stdio_client(server, errlog, process=None):
-  """A wrapper around stdio_client that ensures proper process tracking and cleanup."""
-  our_process = process
-
-  # If no process was provided, create one
-  if our_process is None:
-    our_process = await asyncio.create_subprocess_exec(
-        server.command,
-        *server.args,
-        stdin=asyncio.subprocess.PIPE,
-        stdout=asyncio.subprocess.PIPE,
-        stderr=errlog,
-    )
-
-  # Use the original stdio_client, but ensure process cleanup
-  try:
-    async with stdio_client(server=server, errlog=errlog) as client:
-      yield client, our_process
-  finally:
-    # Ensure the process is properly terminated if it still exists
-    if our_process and our_process.returncode is None:
-      try:
-        logger.info(
-            f'Terminating process {our_process.pid} from tracked_stdio_client'
-        )
-        our_process.terminate()
-        try:
-          await asyncio.wait_for(our_process.wait(), timeout=3.0)
-        except asyncio.TimeoutError:
-          # Force kill if it doesn't terminate quickly
-          if our_process.returncode is None:
-            logger.warning(f'Forcing kill of process {our_process.pid}')
-            our_process.kill()
-      except ProcessLookupError:
-        # Process already gone, that's fine
-        logger.info(f'Process {our_process.pid} already terminated')
-
-
 class MCPSessionManager:
   """Manages MCP client sessions.
 
@@ -166,162 +124,78 @@ class MCPSessionManager:
   def __init__(
       self,
       connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
   ):
     """Initializes the MCP session manager.
 
-    Example usage:
-    ```
-    mcp_session_manager = MCPSessionManager(
-        connection_params=connection_params,
-        exit_stack=exit_stack,
-    )
-    session = await mcp_session_manager.create_session()
-    ```
-
     Args:
         connection_params: Parameters for the MCP connection (Stdio or SSE).
-        exit_stack: AsyncExitStack to manage the session lifecycle.
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
     """
-
     self._connection_params = connection_params
-    self._exit_stack = exit_stack
     self._errlog = errlog
-    self._process = None  # Track the subprocess
-    self._active_processes = set()  # Track all processes created
-    self._active_file_handles = set()  # Track file handles
+    # Each session manager maintains its own exit stack for proper cleanup
+    self._exit_stack: Optional[AsyncExitStack] = None
+    self._session: Optional[ClientSession] = None
 
-  async def create_session(
-      self,
-  ) -> tuple[ClientSession, Optional[asyncio.subprocess.Process]]:
-    """Creates a new MCP session and tracks the associated process."""
-    session, process = await self._initialize_session(
-        connection_params=self._connection_params,
-        exit_stack=self._exit_stack,
-        errlog=self._errlog,
-    )
-    self._process = process  # Store reference to process
-
-    # Track the process
-    if process:
-      self._active_processes.add(process)
-
-    return session, process
-
-  @classmethod
-  async def _initialize_session(
-      cls,
-      *,
-      connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
-      errlog: TextIO = sys.stderr,
-  ) -> tuple[ClientSession, Optional[asyncio.subprocess.Process]]:
-    """Initializes an MCP client session.
-
-    Args:
-        connection_params: Parameters for the MCP connection (Stdio or SSE).
-        exit_stack: AsyncExitStack to manage the session lifecycle.
-        errlog: (Optional) TextIO stream for error logging. Use only for
-          initializing a local stdio MCP session.
+  async def create_session(self) -> ClientSession:
+    """Creates and initializes an MCP client session.
 
     Returns:
         ClientSession: The initialized MCP client session.
     """
-    process = None
-
-    if isinstance(connection_params, StdioServerParameters):
-      # For stdio connections, we need to track the subprocess
-      client, process = await cls._create_stdio_client(
-          server=connection_params,
-          errlog=errlog,
-          exit_stack=exit_stack,
-      )
-    elif isinstance(connection_params, SseServerParams):
-      # For SSE connections, create the client without a subprocess
-      client = sse_client(
-          url=connection_params.url,
-          headers=connection_params.headers,
-          timeout=connection_params.timeout,
-          sse_read_timeout=connection_params.sse_read_timeout,
-      )
-    else:
-      raise ValueError(
-          'Unable to initialize connection. Connection should be'
-          ' StdioServerParameters or SseServerParams, but got'
-          f' {connection_params}'
-      )
+    if self._session is not None:
+      return self._session
 
-    # Create the session with the client
-    transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
-    await session.initialize()
+    # Create a new exit stack for this session
+    self._exit_stack = AsyncExitStack()
 
-    return session, process
-
-  @staticmethod
-  async def _create_stdio_client(
-      server: StdioServerParameters,
-      errlog: TextIO,
-      exit_stack: AsyncExitStack,
-  ) -> tuple[Any, asyncio.subprocess.Process]:
-    """Create stdio client and return both the client and process.
-
-    This implementation adapts to how the MCP stdio_client is created.
-    The actual implementation may need to be adjusted based on the MCP library
-    structure.
-    """
-    # Create the subprocess directly so we can track it
-    process = await asyncio.create_subprocess_exec(
-        server.command,
-        *server.args,
-        stdin=asyncio.subprocess.PIPE,
-        stdout=asyncio.subprocess.PIPE,
-        stderr=errlog,
-    )
-
-    # Create the stdio client using the MCP library
     try:
-      # Method 1: Try using the existing process if stdio_client supports it
-      client = stdio_client(server=server, errlog=errlog, process=process)
-    except TypeError:
-      # Method 2: If the above doesn't work, let stdio_client create its own process
-      # and we'll need to terminate both processes later
-      logger.warning(
-          'Using stdio_client with its own process - may lead to duplicate'
-          ' processes'
-      )
-      client = stdio_client(server=server, errlog=errlog)
+      if isinstance(self._connection_params, StdioServerParameters):
+        client = stdio_client(
+            server=self._connection_params, errlog=self._errlog
+        )
+      elif isinstance(self._connection_params, SseServerParams):
+        client = sse_client(
+            url=self._connection_params.url,
+            headers=self._connection_params.headers,
+            timeout=self._connection_params.timeout,
+            sse_read_timeout=self._connection_params.sse_read_timeout,
+        )
+      else:
+        raise ValueError(
+            'Unable to initialize connection. Connection should be'
+            ' StdioServerParameters or SseServerParams, but got'
+            f' {self._connection_params}'
+        )
 
-    return client, process
+      transports = await self._exit_stack.enter_async_context(client)
+      session = await self._exit_stack.enter_async_context(
+          ClientSession(*transports)
+      )
+      await session.initialize()
 
-  async def _emergency_cleanup(self):
-    """Perform emergency cleanup of resources when normal cleanup fails."""
-    logger.info('Performing emergency cleanup of MCPSessionManager resources')
+      self._session = session
+      return session
 
-    # Clean up any tracked processes
-    for proc in list(self._active_processes):
-      try:
-        if proc and proc.returncode is None:
-          logger.info(f'Emergency termination of process {proc.pid}')
-          proc.terminate()
-          try:
-            await asyncio.wait_for(proc.wait(), timeout=1.0)
-          except asyncio.TimeoutError:
-            logger.warning(f"Process {proc.pid} didn't terminate, forcing kill")
-            proc.kill()
-        self._active_processes.remove(proc)
-      except Exception as e:
-        logger.error(f'Error during process cleanup: {e}')
+    except Exception:
+      # If session creation fails, clean up the exit stack
+      if self._exit_stack:
+        await self._exit_stack.aclose()
+        self._exit_stack = None
+      raise
 
-    # Clean up any tracked file handles
-    for handle in list(self._active_file_handles):
+  async def close(self):
+    """Closes the session and cleans up resources."""
+    if self._exit_stack:
       try:
-        if not handle.closed:
-          logger.info('Closing file handle')
-          handle.close()
-        self._active_file_handles.remove(handle)
+        await self._exit_stack.aclose()
       except Exception as e:
-        logger.error(f'Error closing file handle: {e}')
+        # Log the error but don't re-raise to avoid blocking shutdown
+        print(
+            f'Warning: Error during MCP session cleanup: {e}', file=self._errlog
+        )
+      finally:
+        self._exit_stack = None
+        self._session = None