google-adk 1.5.0__py3-none-any.whl → 1.6.1__py3-none-any.whl

google/adk/runners.py

@@ -307,27 +307,44 @@ class Runner:
     root_agent = self.agent
     invocation_context.agent = self._find_agent_to_run(session, root_agent)
 
+    # Pre-processing for live streaming tools
+    # Inspect the tool's parameters to find if it uses LiveRequestQueue
     invocation_context.active_streaming_tools = {}
     # TODO(hangfei): switch to use canonical_tools.
     # for shell agents, there is no tools associated with it so we should skip.
     if hasattr(invocation_context.agent, 'tools'):
-      for tool in invocation_context.agent.tools:
-        # replicate a LiveRequestQueue for streaming tools that relis on
-        # LiveRequestQueue
-        from typing import get_type_hints
+      import inspect
 
-        type_hints = get_type_hints(tool)
-        for arg_type in type_hints.values():
-          if arg_type is LiveRequestQueue:
+      for tool in invocation_context.agent.tools:
+        # We use `inspect.signature()` to examine the tool's underlying function (`tool.func`).
+        # This approach is deliberately chosen over `typing.get_type_hints()` for robustness.
+        #
+        # The Problem with `get_type_hints()`:
+        # `get_type_hints()` attempts to resolve forward-referenced (string-based) type
+        # annotations. This resolution can easily fail with a `NameError` (e.g., "Union not found")
+        # if the type isn't available in the scope where `get_type_hints()` is called.
+        # This is a common and brittle issue in framework code that inspects functions
+        # defined in separate user modules.
+        #
+        # Why `inspect.signature()` is Better Here:
+        # `inspect.signature()` does NOT resolve the annotations; it retrieves the raw
+        # annotation object as it was defined on the function. This allows us to
+        # perform a direct and reliable identity check (`param.annotation is LiveRequestQueue`)
+        # without risking a `NameError`.
+        callable_to_inspect = tool.func if hasattr(tool, 'func') else tool
+        # Ensure the target is actually callable before inspecting to avoid errors.
+        if not callable(callable_to_inspect):
+          continue
+        for param in inspect.signature(callable_to_inspect).parameters.values():
+          if param.annotation is LiveRequestQueue:
             if not invocation_context.active_streaming_tools:
               invocation_context.active_streaming_tools = {}
-            active_streaming_tools = ActiveStreamingTool(
+            active_streaming_tool = ActiveStreamingTool(
                 stream=LiveRequestQueue()
             )
             invocation_context.active_streaming_tools[tool.__name__] = (
-                active_streaming_tools
+                active_streaming_tool
             )
-
     async for event in invocation_context.agent.run_live(invocation_context):
       await self.session_service.append_event(session=session, event=event)
       yield event

google/adk/sessions/database_session_service.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import copy
 from datetime import datetime
+from datetime import timezone
 import json
 import logging
 from typing import Any
@@ -144,6 +145,21 @@ class StorageSession(Base):
   def __repr__(self):
     return f"<StorageSession(id={self.id}, update_time={self.update_time})>"
 
+  @property
+  def _dialect_name(self) -> Optional[str]:
+    session = inspect(self).session
+    return session.bind.dialect.name if session else None
+
+  @property
+  def update_timestamp_tz(self) -> datetime:
+    """Returns the time zone aware update timestamp."""
+    if self._dialect_name == "sqlite":
+      # SQLite does not support timezone. SQLAlchemy returns a naive datetime
+      # object without timezone information. We need to convert it to UTC
+      # manually.
+      return self.update_time.replace(tzinfo=timezone.utc).timestamp()
+    return self.update_time.timestamp()
+
 
 class StorageEvent(Base):
   """Represents an event stored in the database."""
@@ -412,7 +428,7 @@ class DatabaseSessionService(BaseSessionService):
           user_id=str(storage_session.user_id),
           id=str(storage_session.id),
           state=merged_state,
-          last_update_time=storage_session.update_time.timestamp(),
+          last_update_time=storage_session.update_timestamp_tz,
       )
       return session
 
@@ -443,7 +459,9 @@ class DatabaseSessionService(BaseSessionService):
 
       storage_events = (
           session_factory.query(StorageEvent)
+          .filter(StorageEvent.app_name == app_name)
           .filter(StorageEvent.session_id == storage_session.id)
+          .filter(StorageEvent.user_id == user_id)
           .filter(timestamp_filter)
           .order_by(StorageEvent.timestamp.desc())
           .limit(
@@ -473,7 +491,7 @@ class DatabaseSessionService(BaseSessionService):
           user_id=user_id,
           id=session_id,
           state=merged_state,
-          last_update_time=storage_session.update_time.timestamp(),
+          last_update_time=storage_session.update_timestamp_tz,
       )
       session.events = [e.to_event() for e in reversed(storage_events)]
     return session
@@ -496,7 +514,7 @@ class DatabaseSessionService(BaseSessionService):
             user_id=user_id,
             id=storage_session.id,
             state={},
-            last_update_time=storage_session.update_time.timestamp(),
+            last_update_time=storage_session.update_timestamp_tz,
         )
         sessions.append(session)
       return ListSessionsResponse(sessions=sessions)
@@ -529,13 +547,13 @@ class DatabaseSessionService(BaseSessionService):
           StorageSession, (session.app_name, session.user_id, session.id)
       )
 
-      if storage_session.update_time.timestamp() > session.last_update_time:
+      if storage_session.update_timestamp_tz > session.last_update_time:
         raise ValueError(
             "The last_update_time provided in the session object"
             f" {datetime.fromtimestamp(session.last_update_time):'%Y-%m-%d %H:%M:%S'} is"
             " earlier than the update_time in the storage_session"
-            f" {storage_session.update_time:'%Y-%m-%d %H:%M:%S'}. Please check"
-            " if it is a stale session."
+            f" {datetime.fromtimestamp(storage_session.update_timestamp_tz):'%Y-%m-%d %H:%M:%S'}."
+            " Please check if it is a stale session."
         )
 
       # Fetch states from storage
@@ -577,7 +595,7 @@ class DatabaseSessionService(BaseSessionService):
       session_factory.refresh(storage_session)
 
       # Update timestamp with commit time
-      session.last_update_time = storage_session.update_time.timestamp()
+      session.last_update_time = storage_session.update_timestamp_tz
 
     # Also update the in-memory session
     await super().append_event(session=session, event=event)

google/adk/sessions/in_memory_session_service.py

@@ -33,7 +33,11 @@ logger = logging.getLogger('google_adk.' + __name__)
 
 
 class InMemorySessionService(BaseSessionService):
-  """An in-memory implementation of the session service."""
+  """An in-memory implementation of the session service.
+
+  It is not suitable for multi-threaded production environments. Use it for
+  testing and development only.
+  """
 
   def __init__(self):
     # A map from app name to a map from user ID to a map from session ID to

google/adk/sessions/vertex_ai_session_service.py

@@ -25,6 +25,11 @@ import urllib.parse
 
 from dateutil import parser
 from google.genai.errors import ClientError
+from tenacity import retry
+from tenacity import retry_if_result
+from tenacity import RetryError
+from tenacity import stop_after_attempt
+from tenacity import wait_exponential
 from typing_extensions import override
 
 from google import genai
@@ -64,6 +69,20 @@ class VertexAiSessionService(BaseSessionService):
     self._location = location
     self._agent_engine_id = agent_engine_id
 
+  async def _get_session_api_response(
+      self,
+      reasoning_engine_id: str,
+      session_id: str,
+      api_client: genai.ApiClient,
+  ):
+    get_session_api_response = await api_client.async_request(
+        http_method='GET',
+        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
+        request_dict={},
+    )
+    get_session_api_response = _convert_api_response(get_session_api_response)
+    return get_session_api_response
+
   @override
   async def create_session(
       self,
@@ -95,66 +114,68 @@ class VertexAiSessionService(BaseSessionService):
 
     session_id = api_response['name'].split('/')[-3]
     operation_id = api_response['name'].split('/')[-1]
-
-    max_retry_attempt = 5
-
     if _is_vertex_express_mode(self._project, self._location):
       # Express mode doesn't support LRO, so we need to poll
       # the session resource.
       # TODO: remove this once LRO polling is supported in Express mode.
-      for i in range(max_retry_attempt):
+      @retry(
+          stop=stop_after_attempt(5),
+          wait=wait_exponential(multiplier=1, min=1, max=3),
+          retry=retry_if_result(lambda response: not response),
+          reraise=True,
+      )
+      async def _poll_session_resource():
         try:
-          await api_client.async_request(
-              http_method='GET',
-              path=(
-                  f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
-              ),
-              request_dict={},
+          return await self._get_session_api_response(
+              reasoning_engine_id, session_id, api_client
           )
-          break
-        except ClientError as e:
-          logger.info('Polling for session %s: %s', session_id, e)
-          # Add slight exponential backoff to avoid excessive polling.
-          await asyncio.sleep(1 + 0.5 * i)
-      else:
-        raise TimeoutError('Session creation failed.')
+        except ClientError:
+          logger.info(f'Polling session resource')
+          return None
+
+      try:
+        await _poll_session_resource()
+      except Exception as exc:
+        raise ValueError('Failed to create session.') from exc
     else:
-      lro_response = None
-      for _ in range(max_retry_attempt):
+
+      @retry(
+          stop=stop_after_attempt(5),
+          wait=wait_exponential(multiplier=1, min=1, max=3),
+          retry=retry_if_result(
+              lambda response: not response.get('done', False),
+          ),
+          reraise=True,
+      )
+      async def _poll_lro():
         lro_response = await api_client.async_request(
             http_method='GET',
             path=f'operations/{operation_id}',
             request_dict={},
         )
         lro_response = _convert_api_response(lro_response)
+        return lro_response
 
-        if lro_response.get('done', None):
-          break
-
-        await asyncio.sleep(1)
-
-      if lro_response is None or not lro_response.get('done', None):
+      try:
+        await _poll_lro()
+      except RetryError as exc:
         raise TimeoutError(
             f'Timeout waiting for operation {operation_id} to complete.'
-        )
+        ) from exc
+      except Exception as exc:
+        raise ValueError('Failed to create session.') from exc
 
-    # Get session resource
-    get_session_api_response = await api_client.async_request(
-        http_method='GET',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
-        request_dict={},
+    get_session_api_response = await self._get_session_api_response(
+        reasoning_engine_id, session_id, api_client
     )
-    get_session_api_response = _convert_api_response(get_session_api_response)
-
-    update_timestamp = isoparse(
-        get_session_api_response['updateTime']
-    ).timestamp()
     session = Session(
         app_name=str(app_name),
         user_id=str(user_id),
         id=str(session_id),
         state=get_session_api_response.get('sessionState', {}),
-        last_update_time=update_timestamp,
+        last_update_time=isoparse(
+            get_session_api_response['updateTime']
+        ).timestamp(),
     )
     return session
 
@@ -171,12 +192,12 @@ class VertexAiSessionService(BaseSessionService):
     api_client = self._get_api_client()
 
     # Get session resource
-    get_session_api_response = await api_client.async_request(
-        http_method='GET',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
-        request_dict={},
+    get_session_api_response = await self._get_session_api_response(
+        reasoning_engine_id, session_id, api_client
     )
-    get_session_api_response = _convert_api_response(get_session_api_response)
+
+    if get_session_api_response['userId'] != user_id:
+      raise ValueError(f'Session not found: {session_id}')
 
     session_id = get_session_api_response['name'].split('/')[-1]
     update_timestamp = isoparse(
@@ -360,6 +381,7 @@ def _convert_event_to_json(event: Event) -> Dict[str, Any]:
       'turn_complete': event.turn_complete,
       'interrupted': event.interrupted,
       'branch': event.branch,
+      'custom_metadata': event.custom_metadata,
       'long_running_tool_ids': (
           list(event.long_running_tool_ids)
           if event.long_running_tool_ids
@@ -439,6 +461,9 @@ def _from_api_event(api_event: Dict[str, Any]) -> Event:
     event.turn_complete = api_event['eventMetadata'].get('turnComplete', None)
     event.interrupted = api_event['eventMetadata'].get('interrupted', None)
     event.branch = api_event['eventMetadata'].get('branch', None)
+    event.custom_metadata = api_event['eventMetadata'].get(
+        'customMetadata', None
+    )
     event.grounding_metadata = _session_util.decode_grounding_metadata(
         api_event['eventMetadata'].get('groundingMetadata', None)
     )

google/adk/tools/bigquery/config.py

@@ -26,10 +26,20 @@ class WriteMode(Enum):
 
   BLOCKED = 'blocked'
   """No write operations are allowed.
-  
+
   This mode implies that only read (i.e. SELECT query) operations are allowed.
   """
 
+  PROTECTED = 'protected'
+  """Only protected write operations are allowed in a BigQuery session.
+
+  In this mode write operations in the anonymous dataset of a BigQuery session
+  are allowed. For example, a temporaray table can be created, manipulated and
+  deleted in the anonymous dataset during Agent interaction, while protecting
+  permanent tables from being modified or deleted. To learn more about BigQuery
+  sessions, see https://cloud.google.com/bigquery/docs/sessions-intro.
+  """
+
   ALLOWED = 'allowed'
   """All write operations are allowed."""