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

google/adk/runners.py

@@ -21,8 +21,8 @@ import threading
 from typing import AsyncGenerator
 from typing import Generator
 from typing import Optional
+import warnings
 
-from deprecated import deprecated
 from google.genai import types
 
 from .agents.active_streaming_tool import ActiveStreamingTool
@@ -32,9 +32,9 @@ from .agents.invocation_context import new_invocation_context_id
 from .agents.live_request_queue import LiveRequestQueue
 from .agents.llm_agent import LlmAgent
 from .agents.run_config import RunConfig
-from .agents.run_config import StreamingMode
 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
@@ -42,9 +42,9 @@ 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(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class Runner:
@@ -172,7 +172,7 @@ class Runner:
       The events generated by the agent.
     """
     with tracer.start_as_current_span('invocation'):
-      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:
@@ -196,7 +196,7 @@ class Runner:
       invocation_context.agent = self._find_agent_to_run(session, root_agent)
       async for event in invocation_context.agent.run_async(invocation_context):
         if not event.partial:
-          self.session_service.append_event(session=session, event=event)
+          await self.session_service.append_event(session=session, event=event)
         yield event
 
   async def _append_new_message_to_session(
@@ -241,30 +241,57 @@ class Runner:
         author='user',
         content=new_message,
     )
-    self.session_service.append_event(session=session, event=event)
+    await self.session_service.append_event(session=session, event=event)
 
   async def run_live(
       self,
       *,
-      session: Session,
+      user_id: Optional[str] = None,
+      session_id: Optional[str] = None,
       live_request_queue: LiveRequestQueue,
       run_config: RunConfig = RunConfig(),
+      session: Optional[Session] = None,
   ) -> AsyncGenerator[Event, None]:
     """Runs the agent in live mode (experimental feature).
 
     Args:
-        session: The session to use.
+        user_id: The user ID for the session. Required if `session` is None.
+        session_id: The session ID for the session. Required if `session` is
+          None.
         live_request_queue: The queue for live requests.
         run_config: The run config for the agent.
+        session: The session to use. This parameter is deprecated, please use
+          `user_id` and `session_id` instead.
 
     Yields:
-        The events generated by the agent.
+        AsyncGenerator[Event, None]: An asynchronous generator that yields
+        `Event`
+        objects as they are produced by the agent during its live execution.
 
     .. warning::
         This feature is **experimental** and its API or behavior may change
         in future releases.
+
+    .. note::
+        Either `session` or both `user_id` and `session_id` must be provided.
     """
-    # TODO: right now, only works for a single audio agent without FC.
+    if session is None and (user_id is None or session_id is None):
+      raise ValueError(
+          'Either session or user_id and session_id must be provided.'
+      )
+    if session is not None:
+      warnings.warn(
+          'The `session` parameter is deprecated. Please use `user_id` and'
+          ' `session_id` instead.',
+          DeprecationWarning,
+          stacklevel=2,
+      )
+    if not session:
+      session = await self.session_service.get_session(
+          app_name=self.app_name, user_id=user_id, session_id=session_id
+      )
+      if not session:
+        raise ValueError(f'Session not found: {session_id}')
     invocation_context = self._new_invocation_context_for_live(
         session,
         live_request_queue=live_request_queue,
@@ -276,37 +303,29 @@ class Runner:
 
     invocation_context.active_streaming_tools = {}
     # TODO(hangfei): switch to use canonical_tools.
-    for tool in invocation_context.agent.tools:
-      # replicate a LiveRequestQueue for streaming tools that relis on
-      # LiveRequestQueue
-      from typing import get_type_hints
-
-      type_hints = get_type_hints(tool)
-      for arg_type in type_hints.values():
-        if arg_type is LiveRequestQueue:
-          if not invocation_context.active_streaming_tools:
-            invocation_context.active_streaming_tools = {}
-          active_streaming_tools = ActiveStreamingTool(
-              stream=LiveRequestQueue()
-          )
-          invocation_context.active_streaming_tools[tool.__name__] = (
-              active_streaming_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
+
+        type_hints = get_type_hints(tool)
+        for arg_type in type_hints.values():
+          if arg_type is LiveRequestQueue:
+            if not invocation_context.active_streaming_tools:
+              invocation_context.active_streaming_tools = {}
+            active_streaming_tools = ActiveStreamingTool(
+                stream=LiveRequestQueue()
+            )
+            invocation_context.active_streaming_tools[tool.__name__] = (
+                active_streaming_tools
+            )
 
     async for event in invocation_context.agent.run_live(invocation_context):
-      self.session_service.append_event(session=session, event=event)
+      await self.session_service.append_event(session=session, event=event)
       yield event
 
-  async def close_session(self, session: Session):
-    """Closes a session and adds it to the memory service (experimental feature).
-
-    Args:
-        session: The session to close.
-    """
-    if self.memory_service:
-      await self.memory_service.add_session_to_memory(session)
-    self.session_service.close_session(session=session)
-
   def _find_agent_to_run(
       self, session: Session, root_agent: BaseAgent
   ) -> BaseAgent:
@@ -391,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,
@@ -430,12 +449,46 @@ class Runner:
           run_config.output_audio_transcription = (
               types.AudioTranscriptionConfig()
           )
+      if not run_config.input_audio_transcription:
+        # need this input transcription for agent transferring in live mode.
+        run_config.input_audio_transcription = types.AudioTranscriptionConfig()
     return self._new_invocation_context(
         session,
         live_request_queue=live_request_queue,
         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.
@@ -448,9 +501,11 @@ class InMemoryRunner(Runner):
       agent: The root agent to run.
       app_name: The application name of the runner. Defaults to
         'InMemoryRunner'.
+      _in_memory_session_service: Deprecated. Please don't use. The in-memory
+        session service for the runner.
   """
 
-  def __init__(self, agent: LlmAgent, *, app_name: str = 'InMemoryRunner'):
+  def __init__(self, agent: BaseAgent, *, app_name: str = 'InMemoryRunner'):
     """Initializes the InMemoryRunner.
 
     Args:
@@ -458,10 +513,11 @@ class InMemoryRunner(Runner):
         app_name: The application name of the runner. Defaults to
           'InMemoryRunner'.
     """
+    self._in_memory_session_service = InMemorySessionService()
     super().__init__(
         app_name=app_name,
         agent=agent,
         artifact_service=InMemoryArtifactService(),
-        session_service=InMemorySessionService(),
+        session_service=self._in_memory_session_service,
         memory_service=InMemoryMemoryService(),
     )

google/adk/sessions/__init__.py

@@ -19,7 +19,7 @@ from .session import Session
 from .state import State
 from .vertex_ai_session_service import VertexAiSessionService
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 __all__ = [

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/base_session_service.py

@@ -40,13 +40,6 @@ class ListSessionsResponse(BaseModel):
   sessions: list[Session] = Field(default_factory=list)
 
 
-class ListEventsResponse(BaseModel):
-  """The response of listing events in a session."""
-
-  events: list[Event] = Field(default_factory=list)
-  next_page_token: Optional[str] = None
-
-
 class BaseSessionService(abc.ABC):
   """Base class for session services.
 
@@ -54,7 +47,7 @@ class BaseSessionService(abc.ABC):
   """
 
   @abc.abstractmethod
-  def create_session(
+  async def create_session(
       self,
       *,
       app_name: str,
@@ -74,10 +67,9 @@ class BaseSessionService(abc.ABC):
     Returns:
       session: The newly created session instance.
     """
-    pass
 
   @abc.abstractmethod
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
@@ -86,39 +78,20 @@ class BaseSessionService(abc.ABC):
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
     """Gets a session."""
-    pass
 
   @abc.abstractmethod
-  def list_sessions(
+  async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
     """Lists all the sessions."""
-    pass
 
   @abc.abstractmethod
-  def delete_session(
+  async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     """Deletes a session."""
-    pass
-
-  @abc.abstractmethod
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    """Lists events in a session."""
-    pass
-
-  def close_session(self, *, session: Session):
-    """Closes a session."""
-    # TODO: determine whether we want to finalize the session here.
-    pass
 
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     """Appends an event to a session object."""
     if event.partial:
       return event
@@ -126,7 +99,7 @@ class BaseSessionService(abc.ABC):
     session.events.append(event)
     return event
 
-  def __update_session_state(self, session: Session, event: Event):
+  def __update_session_state(self, session: Session, event: Event) -> None:
     """Updates the session state based on the event."""
     if not event.actions or not event.actions.state_delta:
       return

google/adk/sessions/database_session_service.py

@@ -15,7 +15,8 @@ import copy
 from datetime import datetime
 import json
 import logging
-from typing import Any, Optional
+from typing import Any
+from typing import Optional
 import uuid
 
 from sqlalchemy import Boolean
@@ -45,27 +46,22 @@ 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 ListEventsResponse
 from .base_session_service import ListSessionsResponse
 from .session import Session
 from .state import State
 
-
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 DEFAULT_MAX_KEY_LENGTH = 128
 DEFAULT_MAX_VARCHAR_LENGTH = 256
 
 
 class DynamicJSON(TypeDecorator):
-  """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON
-
-  serialization for other databases.
-  """
+  """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON serialization for other databases."""
 
   impl = Text  # Default implementation is TEXT
 
@@ -243,10 +239,7 @@ class DatabaseSessionService(BaseSessionService):
   """A session service that uses a database for storage."""
 
   def __init__(self, db_url: str):
-    """
-    Args:
-        db_url: The database URL to connect to.
-    """
+    """Initializes the database session service with a database URL."""
     # 1. Create DB engine for db connection
     # 2. Create all tables based on schema
     # 3. Initialize all properties
@@ -275,7 +268,7 @@ class DatabaseSessionService(BaseSessionService):
     self.inspector = inspect(self.db_engine)
 
     # DB session factory method
-    self.DatabaseSessionFactory: sessionmaker[DatabaseSessionFactory] = (
+    self.database_session_factory: sessionmaker[DatabaseSessionFactory] = (
         sessionmaker(bind=self.db_engine)
     )
 
@@ -284,7 +277,7 @@ class DatabaseSessionService(BaseSessionService):
     Base.metadata.create_all(self.db_engine)
 
   @override
-  def create_session(
+  async def create_session(
       self,
       *,
       app_name: str,
@@ -298,11 +291,11 @@ class DatabaseSessionService(BaseSessionService):
     # 4. Build the session object with generated id
     # 5. Return the session
 
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
 
       # Fetch app and user states from storage
-      storage_app_state = sessionFactory.get(StorageAppState, (app_name))
-      storage_user_state = sessionFactory.get(
+      storage_app_state = session_factory.get(StorageAppState, (app_name))
+      storage_user_state = session_factory.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -312,12 +305,12 @@ class DatabaseSessionService(BaseSessionService):
       # Create state tables if not exist
       if not storage_app_state:
         storage_app_state = StorageAppState(app_name=app_name, state={})
-        sessionFactory.add(storage_app_state)
+        session_factory.add(storage_app_state)
       if not storage_user_state:
         storage_user_state = StorageUserState(
             app_name=app_name, user_id=user_id, state={}
         )
-        sessionFactory.add(storage_user_state)
+        session_factory.add(storage_user_state)
 
       # Extract state deltas
       app_state_delta, user_state_delta, session_state = _extract_state_delta(
@@ -341,10 +334,10 @@ class DatabaseSessionService(BaseSessionService):
           id=session_id,
           state=session_state,
       )
-      sessionFactory.add(storage_session)
-      sessionFactory.commit()
+      session_factory.add(storage_session)
+      session_factory.commit()
 
-      sessionFactory.refresh(storage_session)
+      session_factory.refresh(storage_session)
 
       # Merge states for response
       merged_state = _merge_state(app_state, user_state, session_state)
@@ -358,7 +351,7 @@ class DatabaseSessionService(BaseSessionService):
       return session
 
   @override
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
@@ -369,29 +362,35 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Get the storage session entry from session table
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
-    with self.DatabaseSessionFactory() as sessionFactory:
-      storage_session = sessionFactory.get(
+    with self.database_session_factory() as session_factory:
+      storage_session = session_factory.get(
           StorageSession, (app_name, user_id, session_id)
       )
       if storage_session is None:
         return None
 
+      if config and config.after_timestamp:
+        after_dt = datetime.fromtimestamp(config.after_timestamp)
+        timestamp_filter = StorageEvent.timestamp >= after_dt
+      else:
+        timestamp_filter = True
+
       storage_events = (
-          sessionFactory.query(StorageEvent)
+          session_factory.query(StorageEvent)
           .filter(StorageEvent.session_id == storage_session.id)
-          .filter(
-              StorageEvent.timestamp < config.after_timestamp
-              if config
-              else True
+          .filter(timestamp_filter)
+          .order_by(StorageEvent.timestamp.desc())
+          .limit(
+              config.num_recent_events
+              if config and config.num_recent_events
+              else None
           )
-          .limit(config.num_recent_events if config else None)
-          .order_by(StorageEvent.timestamp.asc())
           .all()
       )
 
       # Fetch states from storage
-      storage_app_state = sessionFactory.get(StorageAppState, (app_name))
-      storage_user_state = sessionFactory.get(
+      storage_app_state = session_factory.get(StorageAppState, (app_name))
+      storage_user_state = session_factory.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -427,17 +426,17 @@ class DatabaseSessionService(BaseSessionService):
               error_message=e.error_message,
               interrupted=e.interrupted,
           )
-          for e in storage_events
+          for e in reversed(storage_events)
       ]
     return session
 
   @override
-  def list_sessions(
+  async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
       results = (
-          sessionFactory.query(StorageSession)
+          session_factory.query(StorageSession)
           .filter(StorageSession.app_name == app_name)
           .filter(StorageSession.user_id == user_id)
           .all()
@@ -455,20 +454,20 @@ class DatabaseSessionService(BaseSessionService):
       return ListSessionsResponse(sessions=sessions)
 
   @override
-  def delete_session(
+  async def delete_session(
       self, app_name: str, user_id: str, session_id: str
   ) -> None:
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
       stmt = delete(StorageSession).where(
           StorageSession.app_name == app_name,
           StorageSession.user_id == user_id,
           StorageSession.id == session_id,
       )
-      sessionFactory.execute(stmt)
-      sessionFactory.commit()
+      session_factory.execute(stmt)
+      session_factory.commit()
 
   @override
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     logger.info(f"Append event: {event} to session {session.id}")
 
     if event.partial:
@@ -477,24 +476,25 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Check if timestamp is stale
     # 2. Update session attributes based on event config
     # 3. Store event to table
-    with self.DatabaseSessionFactory() as sessionFactory:
-      storage_session = sessionFactory.get(
+    with self.database_session_factory() as session_factory:
+      storage_session = session_factory.get(
           StorageSession, (session.app_name, session.user_id, session.id)
       )
 
       if storage_session.update_time.timestamp() > session.last_update_time:
         raise ValueError(
-          f"Session last_update_time "
-          f"{datetime.fromtimestamp(session.last_update_time):%Y-%m-%d %H:%M:%S} "
-          f"is later than the update_time in storage "
-          f"{storage_session.update_time:%Y-%m-%d %H:%M:%S}"
-      )
+            "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."
+        )
 
       # Fetch states from storage
-      storage_app_state = sessionFactory.get(
+      storage_app_state = session_factory.get(
           StorageAppState, (session.app_name)
       )
-      storage_user_state = sessionFactory.get(
+      storage_user_state = session_factory.get(
           StorageUserState, (session.app_name, session.user_id)
       )
 
@@ -543,28 +543,18 @@ class DatabaseSessionService(BaseSessionService):
       if event.content:
         storage_event.content = _session_util.encode_content(event.content)
 
-      sessionFactory.add(storage_event)
+      session_factory.add(storage_event)
 
-      sessionFactory.commit()
-      sessionFactory.refresh(storage_session)
+      session_factory.commit()
+      session_factory.refresh(storage_session)
 
       # Update timestamp with commit time
       session.last_update_time = storage_session.update_time.timestamp()
 
     # Also update the in-memory session
-    super().append_event(session=session, event=event)
+    await super().append_event(session=session, event=event)
     return event
 
-  @override
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    raise NotImplementedError()
-
 
 def convert_event(event: StorageEvent) -> Event:
   """Converts a storage event to an event."""