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

google/adk/models/google_llm.py

@@ -98,6 +98,7 @@ class Gemini(BaseLlm):
       )
       response = None
       text = ''
+      usage_metadata = None
       # for sse, similar as bidi (see receive method in gemini_llm_connecton.py),
       # we need to mark those text content as partial and after all partial
       # contents are sent, we send an accumulated event which contains all the
@@ -106,6 +107,7 @@ class Gemini(BaseLlm):
       async for response in responses:
         logger.info(_build_response_log(response))
         llm_response = LlmResponse.create(response)
+        usage_metadata = llm_response.usage_metadata
         if (
             llm_response.content
             and llm_response.content.parts
@@ -123,7 +125,7 @@ class Gemini(BaseLlm):
               content=types.ModelContent(
                   parts=[types.Part.from_text(text=text)],
               ),
-              usage_metadata=llm_response.usage_metadata,
+              usage_metadata=usage_metadata,
           )
           text = ''
         yield llm_response
@@ -137,6 +139,7 @@ class Gemini(BaseLlm):
             content=types.ModelContent(
                 parts=[types.Part.from_text(text=text)],
             ),
+            usage_metadata=usage_metadata,
         )
 
     else:

google/adk/models/lite_llm.py

@@ -30,6 +30,7 @@ from typing import Union
 from google.genai import types
 from litellm import acompletion
 from litellm import ChatCompletionAssistantMessage
+from litellm import ChatCompletionAssistantToolCall
 from litellm import ChatCompletionDeveloperMessage
 from litellm import ChatCompletionImageUrlObject
 from litellm import ChatCompletionMessageToolCall
@@ -180,12 +181,12 @@ def _content_to_message_param(
     for part in content.parts:
       if part.function_call:
         tool_calls.append(
-            ChatCompletionMessageToolCall(
+            ChatCompletionAssistantToolCall(
                 type="function",
                 id=part.function_call.id,
                 function=Function(
                     name=part.function_call.name,
-                    arguments=part.function_call.args,
+                    arguments=json.dumps(part.function_call.args),
                 ),
             )
         )

google/adk/models/llm_response.py

@@ -14,7 +14,8 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any
+from typing import Optional
 
 from google.genai import types
 from pydantic import alias_generators

google/adk/runners.py

@@ -34,6 +34,7 @@ from .agents.llm_agent import LlmAgent
 from .agents.run_config import RunConfig
 from .artifacts.base_artifact_service import BaseArtifactService
 from .artifacts.in_memory_artifact_service import InMemoryArtifactService
+from .code_executors.built_in_code_executor import BuiltInCodeExecutor
 from .events.event import Event
 from .memory.base_memory_service import BaseMemoryService
 from .memory.in_memory_memory_service import InMemoryMemoryService
@@ -41,7 +42,7 @@ from .sessions.base_session_service import BaseSessionService
 from .sessions.in_memory_session_service import InMemorySessionService
 from .sessions.session import Session
 from .telemetry import tracer
-from .tools._built_in_code_execution_tool import built_in_code_execution
+from .tools.base_toolset import BaseToolset
 
 logger = logging.getLogger('google_adk.' + __name__)
 
@@ -286,7 +287,7 @@ class Runner:
           stacklevel=2,
       )
     if not session:
-      session = self.session_service.get_session(
+      session = await self.session_service.get_session(
           app_name=self.app_name, user_id=user_id, session_id=session_id
       )
       if not session:
@@ -409,8 +410,8 @@ class Runner:
             f'CFC is not supported for model: {model_name} in agent:'
             f' {self.agent.name}'
         )
-      if built_in_code_execution not in self.agent.canonical_tools():
-        self.agent.tools.append(built_in_code_execution)
+      if not isinstance(self.agent.code_executor, BuiltInCodeExecutor):
+        self.agent.code_executor = BuiltInCodeExecutor()
 
     return InvocationContext(
         artifact_service=self.artifact_service,
@@ -457,6 +458,37 @@ class Runner:
         run_config=run_config,
     )
 
+  def _collect_toolset(self, agent: BaseAgent) -> set[BaseToolset]:
+    toolsets = set()
+    if isinstance(agent, LlmAgent):
+      for tool_union in agent.tools:
+        if isinstance(tool_union, BaseToolset):
+          toolsets.add(tool_union)
+    for sub_agent in agent.sub_agents:
+      toolsets.update(self._collect_toolset(sub_agent))
+    return toolsets
+
+  async def _cleanup_toolsets(self, toolsets_to_close: set[BaseToolset]):
+    """Clean up toolsets with proper task context management."""
+    if not toolsets_to_close:
+      return
+
+    # This maintains the same task context throughout cleanup
+    for toolset in toolsets_to_close:
+      try:
+        logger.info('Closing toolset: %s', type(toolset).__name__)
+        # Use asyncio.wait_for to add timeout protection
+        await asyncio.wait_for(toolset.close(), timeout=10.0)
+        logger.info('Successfully closed toolset: %s', type(toolset).__name__)
+      except asyncio.TimeoutError:
+        logger.warning('Toolset %s cleanup timed out', type(toolset).__name__)
+      except Exception as e:
+        logger.error('Error closing toolset %s: %s', type(toolset).__name__, e)
+
+  async def close(self):
+    """Closes the runner."""
+    await self._cleanup_toolsets(self._collect_toolset(self.agent))
+
 
 class InMemoryRunner(Runner):
   """An in-memory Runner for testing and development.

google/adk/sessions/_session_util.py

@@ -15,7 +15,8 @@
 """Utility functions for session service."""
 
 import base64
-from typing import Any, Optional
+from typing import Any
+from typing import Optional
 
 from google.genai import types
 

google/adk/sessions/database_session_service.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import copy
 from datetime import datetime
-from datetime import timezone
 import json
 import logging
 from typing import Any
@@ -47,8 +46,8 @@ from sqlalchemy.types import TypeDecorator
 from typing_extensions import override
 from tzlocal import get_localzone
 
-from ..events.event import Event
 from . import _session_util
+from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
 from .base_session_service import ListSessionsResponse
@@ -371,10 +370,8 @@ class DatabaseSessionService(BaseSessionService):
         return None
 
       if config and config.after_timestamp:
-        after_dt = datetime.fromtimestamp(
-            config.after_timestamp, tz=timezone.utc
-        )
-        timestamp_filter = StorageEvent.timestamp > after_dt
+        after_dt = datetime.fromtimestamp(config.after_timestamp)
+        timestamp_filter = StorageEvent.timestamp >= after_dt
       else:
         timestamp_filter = True
 
@@ -382,7 +379,7 @@ class DatabaseSessionService(BaseSessionService):
           session_factory.query(StorageEvent)
           .filter(StorageEvent.session_id == storage_session.id)
           .filter(timestamp_filter)
-          .order_by(StorageEvent.timestamp.asc())
+          .order_by(StorageEvent.timestamp.desc())
           .limit(
               config.num_recent_events
               if config and config.num_recent_events
@@ -429,7 +426,7 @@ class DatabaseSessionService(BaseSessionService):
               error_message=e.error_message,
               interrupted=e.interrupted,
           )
-          for e in storage_events
+          for e in reversed(storage_events)
       ]
     return session
 

google/adk/sessions/vertex_ai_session_service.py

@@ -11,19 +11,20 @@
 # 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.
+import asyncio
 import logging
 import re
-import time
 from typing import Any
 from typing import Optional
 
 from dateutil import parser
-from google import genai
 from typing_extensions import override
 
+from google import genai
+
+from . import _session_util
 from ..events.event import Event
 from ..events.event_actions import EventActions
-from . import _session_util
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
 from .base_session_service import ListSessionsResponse
@@ -68,7 +69,8 @@ class VertexAiSessionService(BaseSessionService):
     if state:
       session_json_dict['session_state'] = state
 
-    api_response = await self.api_client.async_request(
+    api_client = _get_api_client(self.project, self.location)
+    api_response = await api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions',
         request_dict=session_json_dict,
@@ -80,7 +82,7 @@ class VertexAiSessionService(BaseSessionService):
 
     max_retry_attempt = 5
     while max_retry_attempt >= 0:
-      lro_response = await self.api_client.async_request(
+      lro_response = await api_client.async_request(
           http_method='GET',
           path=f'operations/{operation_id}',
           request_dict={},
@@ -89,11 +91,11 @@ class VertexAiSessionService(BaseSessionService):
       if lro_response.get('done', None):
         break
 
-      time.sleep(1)
+      await asyncio.sleep(1)
       max_retry_attempt -= 1
 
     # Get session resource
-    get_session_api_response = await self.api_client.async_request(
+    get_session_api_response = await api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
@@ -123,7 +125,8 @@ class VertexAiSessionService(BaseSessionService):
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
 
     # Get session resource
-    get_session_api_response = await self.api_client.async_request(
+    api_client = _get_api_client(self.project, self.location)
+    get_session_api_response = await api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
@@ -141,7 +144,7 @@ class VertexAiSessionService(BaseSessionService):
         last_update_time=update_timestamp,
     )
 
-    list_events_api_response = await self.api_client.async_request(
+    list_events_api_response = await api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events',
         request_dict={},
@@ -180,7 +183,8 @@ class VertexAiSessionService(BaseSessionService):
   ) -> ListSessionsResponse:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
 
-    api_response = self.api_client.request(
+    api_client = _get_api_client(self.project, self.location)
+    api_response = await api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions?filter=user_id={user_id}',
         request_dict={},
@@ -206,7 +210,8 @@ class VertexAiSessionService(BaseSessionService):
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
-    await self.api_client.async_request(
+    api_client = _get_api_client(self.project, self.location)
+    await api_client.async_request(
         http_method='DELETE',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
@@ -218,15 +223,25 @@ class VertexAiSessionService(BaseSessionService):
     await super().append_event(session=session, event=event)
 
     reasoning_engine_id = _parse_reasoning_engine_id(session.app_name)
-    await self.api_client.async_request(
+    api_client = _get_api_client(self.project, self.location)
+    await api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
         request_dict=_convert_event_to_json(event),
     )
-
     return event
 
 
+def _get_api_client(project: str, location: str):
+  """Instantiates an API client for the given project and location.
+
+  It needs to be instantiated inside each request so that the event loop
+  management.
+  """
+  client = genai.Client(vertexai=True, project=project, location=location)
+  return client._api_client
+
+
 def _convert_event_to_json(event: Event):
   metadata_json = {
       'partial': event.partial,

google/adk/telemetry.py

@@ -32,7 +32,6 @@ from .events.event import Event
 from .models.llm_request import LlmRequest
 from .models.llm_response import LlmResponse
 
-
 tracer = trace.get_tracer('gcp.vertex.agent')
 
 
@@ -118,7 +117,10 @@ def trace_call_llm(
   # Consider removing once GenAI SDK provides a way to record this info.
   span.set_attribute(
       'gcp.vertex.agent.llm_request',
-      json.dumps(_build_llm_request_for_trace(llm_request)),
+      json.dumps(
+          _build_llm_request_for_trace(llm_request),
+          default=lambda o: '<not serializable>',
+      ),
   )
   # Consider removing once GenAI SDK provides a way to record this info.
   span.set_attribute(

google/adk/tools/agent_tool.py

@@ -21,10 +21,10 @@ from google.genai import types
 from pydantic import model_validator
 from typing_extensions import override
 
+from . import _automatic_function_calling_util
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
 from ..sessions.in_memory_session_service import InMemorySessionService
-from . import _automatic_function_calling_util
 from .base_tool import BaseTool
 from .tool_context import ToolContext
 

google/adk/tools/apihub_tool/apihub_toolset.py

@@ -131,6 +131,7 @@ class APIHubToolset(BaseToolset):
           be either a tool predicate or a list of tool names of the tools to
           expose.
     """
+    super().__init__(tool_filter=tool_filter)
     self.name = name
     self.description = description
     self._apihub_resource_name = apihub_resource_name
@@ -143,7 +144,6 @@ class APIHubToolset(BaseToolset):
     self._openapi_toolset = None
     self._auth_scheme = auth_scheme
     self._auth_credential = auth_credential
-    self.tool_filter = tool_filter
 
     if not self._lazy_load_spec:
       self._prepare_toolset()

google/adk/tools/apihub_tool/clients/apihub_client.py

@@ -12,11 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
 import base64
 import json
-from typing import Any, Dict, List, Optional, Tuple
-from urllib.parse import parse_qs, urlparse
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from urllib.parse import parse_qs
+from urllib.parse import urlparse
+
 from google.auth import default as default_service_credential
 from google.auth.transport.requests import Request
 from google.oauth2 import service_account

google/adk/tools/apihub_tool/clients/secret_client.py

@@ -14,6 +14,7 @@
 
 import json
 from typing import Optional
+
 import google.auth
 from google.auth import default as default_service_credential
 import google.auth.transport.requests

google/adk/tools/application_integration_tool/application_integration_toolset.py

@@ -128,6 +128,7 @@ class ApplicationIntegrationToolset(BaseToolset):
         Exception: If there is an error during the initialization of the
             integration or connection client.
     """
+    super().__init__(tool_filter=tool_filter)
     self.project = project
     self.location = location
     self._integration = integration
@@ -140,7 +141,6 @@ class ApplicationIntegrationToolset(BaseToolset):
     self._service_account_json = service_account_json
     self._auth_scheme = auth_scheme
     self._auth_credential = auth_credential
-    self.tool_filter = tool_filter
 
     integration_client = IntegrationClient(
         project,
@@ -263,7 +263,11 @@ class ApplicationIntegrationToolset(BaseToolset):
       readonly_context: Optional[ReadonlyContext] = None,
   ) -> List[RestApiTool]:
     return (
-        self._tools
+        [
+            tool
+            for tool in self._tools
+            if self._is_tool_selected(tool, readonly_context)
+        ]
         if self._openapi_toolset is None
         else await self._openapi_toolset.get_tools(readonly_context)
     )

google/adk/tools/application_integration_tool/clients/connections_client.py

@@ -14,7 +14,11 @@
 
 import json
 import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
 
 import google.auth
 from google.auth import default as default_service_credential
@@ -728,6 +732,9 @@ class ConnectionsClient:
             "query": {"$ref": "#/components/schemas/query"},
             "timeout": {"$ref": "#/components/schemas/timeout"},
             "pageSize": {"$ref": "#/components/schemas/pageSize"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 

google/adk/tools/application_integration_tool/clients/integration_client.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import json
-from typing import List, Optional
+from typing import List
+from typing import Optional
+
 from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
 import google.auth
 from google.auth import default as default_service_credential

google/adk/tools/application_integration_tool/integration_connector_tool.py

@@ -21,9 +21,9 @@ from typing import Union
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
+from .. import BaseTool
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
-from .. import BaseTool
 from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
 from ..openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
 from ..openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler

google/adk/tools/base_toolset.py

@@ -1,7 +1,25 @@
+# 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 abc import ABC
 from abc import abstractmethod
-from typing import Optional, runtime_checkable
+from typing import List
+from typing import Optional
 from typing import Protocol
+from typing import runtime_checkable
+from typing import Union
 
 from ..agents.readonly_context import ReadonlyContext
 from .base_tool import BaseTool
@@ -33,9 +51,15 @@ class BaseToolset(ABC):
   A toolset is a collection of tools that can be used by an agent.
   """
 
+  def __init__(
+      self, *, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None
+  ):
+    self.tool_filter = tool_filter
+
   @abstractmethod
   async def get_tools(
-      self, readonly_context: Optional[ReadonlyContext] = None
+      self,
+      readonly_context: Optional[ReadonlyContext] = None,
   ) -> list[BaseTool]:
     """Return all tools in the toolset based on the provided context.
 
@@ -56,3 +80,17 @@ class BaseToolset(ABC):
     should ensure that any open connections, files, or other managed
     resources are properly released to prevent leaks.
     """
+
+  def _is_tool_selected(
+      self, tool: BaseTool, 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

google/adk/tools/bigquery/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+"""BigQuery Tools. (Experimental)
+
+BigQuery Tools under this module are hand crafted and customized while the tools
+under google.adk.tools.google_api_tool are auto generated based on API
+definition. The rationales to have customized tool are:
+
+1. BigQuery APIs have functions overlaps and LLM can't tell what tool to use
+2. BigQuery APIs have a lot of parameters with some rarely used, which are not
+ LLM-friendly
+3. We want to provide more high-level tools like forecasting, RAG, segmentation,
+ etc.
+4. We want to provide extra access guardrails in those tools. For example,
+ execute_sql can't arbitrarily mutate existing data.
+"""