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

google/adk/runners.py

@@ -17,9 +17,14 @@ from __future__ import annotations
 import asyncio
 import logging
 import queue
+import time
+from typing import Any
 from typing import AsyncGenerator
+from typing import Callable
 from typing import Generator
+from typing import List
 from typing import Optional
+import uuid
 import warnings
 
 from google.genai import types
@@ -36,10 +41,13 @@ from .artifacts.in_memory_artifact_service import InMemoryArtifactService
 from .auth.credential_service.base_credential_service import BaseCredentialService
 from .code_executors.built_in_code_executor import BuiltInCodeExecutor
 from .events.event import Event
+from .events.event import EventActions
 from .flows.llm_flows.functions import find_matching_function_call
 from .memory.base_memory_service import BaseMemoryService
 from .memory.in_memory_memory_service import InMemoryMemoryService
 from .platform.thread import create_thread
+from .plugins.base_plugin import BasePlugin
+from .plugins.plugin_manager import PluginManager
 from .sessions.base_session_service import BaseSessionService
 from .sessions.in_memory_session_service import InMemorySessionService
 from .sessions.session import Session
@@ -60,6 +68,7 @@ class Runner:
       app_name: The application name of the runner.
       agent: The root agent to run.
       artifact_service: The artifact service for the runner.
+      plugin_manager: The plugin manager for the runner.
       session_service: The session service for the runner.
       memory_service: The memory service for the runner.
   """
@@ -70,6 +79,8 @@ class Runner:
   """The root agent to run."""
   artifact_service: Optional[BaseArtifactService] = None
   """The artifact service for the runner."""
+  plugin_manager: PluginManager
+  """The plugin manager for the runner."""
   session_service: BaseSessionService
   """The session service for the runner."""
   memory_service: Optional[BaseMemoryService] = None
@@ -82,6 +93,7 @@ class Runner:
       *,
       app_name: str,
       agent: BaseAgent,
+      plugins: Optional[List[BasePlugin]] = None,
       artifact_service: Optional[BaseArtifactService] = None,
       session_service: BaseSessionService,
       memory_service: Optional[BaseMemoryService] = None,
@@ -102,6 +114,7 @@ class Runner:
     self.session_service = session_service
     self.memory_service = memory_service
     self.credential_service = credential_service
+    self.plugin_manager = PluginManager(plugins=plugins)
 
   def run(
       self,
@@ -113,8 +126,9 @@ class Runner:
   ) -> Generator[Event, None, None]:
     """Runs the agent.
 
-    NOTE: This sync interface is only for local testing and convenience purpose.
-    Consider using `run_async` for production usage.
+    NOTE:
+      This sync interface is only for local testing and convenience purpose.
+      Consider using `run_async` for production usage.
 
     Args:
       user_id: The user ID of the session.
@@ -164,6 +178,7 @@ class Runner:
       user_id: str,
       session_id: str,
       new_message: types.Content,
+      state_delta: Optional[dict[str, Any]] = None,
       run_config: RunConfig = RunConfig(),
   ) -> AsyncGenerator[Event, None]:
     """Main entry method to run the agent in this runner.
@@ -191,19 +206,83 @@ class Runner:
       )
       root_agent = self.agent
 
+      # Modify user message before execution.
+      modified_user_message = (
+          await invocation_context.plugin_manager.run_on_user_message_callback(
+              invocation_context=invocation_context, user_message=new_message
+          )
+      )
+      if modified_user_message is not None:
+        new_message = modified_user_message
+
       if new_message:
         await self._append_new_message_to_session(
             session,
             new_message,
             invocation_context,
             run_config.save_input_blobs_as_artifacts,
+            state_delta,
         )
 
       invocation_context.agent = self._find_agent_to_run(session, root_agent)
-      async for event in invocation_context.agent.run_async(invocation_context):
+
+      async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
+        async for event in ctx.agent.run_async(ctx):
+          yield event
+
+      async for event in self._exec_with_plugin(
+          invocation_context, session, execute
+      ):
+        yield event
+
+  async def _exec_with_plugin(
+      self,
+      invocation_context: InvocationContext,
+      session: Session,
+      execute_fn: Callable[[InvocationContext], AsyncGenerator[Event, None]],
+  ) -> AsyncGenerator[Event, None]:
+    """Wraps execution with plugin callbacks.
+
+    Args:
+      invocation_context: The invocation context
+      session: The current session
+      execute_fn: A callable that returns an AsyncGenerator of Events
+
+    Yields:
+      Events from the execution, including any generated by plugins
+    """
+
+    plugin_manager = invocation_context.plugin_manager
+
+    # Step 1: Run the before_run callbacks to see if we should early exit.
+    early_exit_result = await plugin_manager.run_before_run_callback(
+        invocation_context=invocation_context
+    )
+    if isinstance(early_exit_result, Event):
+      await self.session_service.append_event(
+          session=session,
+          event=Event(
+              invocation_id=invocation_context.invocation_id,
+              author='model',
+              content=early_exit_result,
+          ),
+      )
+      yield early_exit_result
+    else:
+      # Step 2: Otherwise continue with normal execution
+      async for event in execute_fn(invocation_context):
         if not event.partial:
           await self.session_service.append_event(session=session, event=event)
-        yield event
+        # Step 3: Run the on_event callbacks to optionally modify the event.
+        modified_event = await plugin_manager.run_on_event_callback(
+            invocation_context=invocation_context, event=event
+        )
+        yield (modified_event if modified_event else event)
+
+    # Step 4: Run the after_run callbacks to optionally modify the context.
+    await plugin_manager.run_after_run_callback(
+        invocation_context=invocation_context
+    )
 
   async def _append_new_message_to_session(
       self,
@@ -211,6 +290,7 @@ class Runner:
       new_message: types.Content,
       invocation_context: InvocationContext,
       save_input_blobs_as_artifacts: bool = False,
+      state_delta: Optional[dict[str, Any]] = None,
   ):
     """Appends a new message to the session.
 
@@ -242,11 +322,19 @@ class Runner:
             text=f'Uploaded file: {file_name}. It is saved into artifacts'
         )
     # Appends only. We do not yield the event because it's not from the model.
-    event = Event(
-        invocation_id=invocation_context.invocation_id,
-        author='user',
-        content=new_message,
-    )
+    if state_delta:
+      event = Event(
+          invocation_id=invocation_context.invocation_id,
+          author='user',
+          actions=EventActions(state_delta=state_delta),
+          content=new_message,
+      )
+    else:
+      event = Event(
+          invocation_id=invocation_context.invocation_id,
+          author='user',
+          content=new_message,
+      )
     await self.session_service.append_event(session=session, event=event)
 
   async def run_live(
@@ -278,7 +366,7 @@ class Runner:
         This feature is **experimental** and its API or behavior may change
         in future releases.
 
-    .. note::
+    .. NOTE::
         Either `session` or both `user_id` and `session_id` must be provided.
     """
     if session is None and (user_id is None or session_id is None):
@@ -345,8 +433,14 @@ class Runner:
             invocation_context.active_streaming_tools[tool.__name__] = (
                 active_streaming_tool
             )
-    async for event in invocation_context.agent.run_live(invocation_context):
-      await self.session_service.append_event(session=session, event=event)
+
+    async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
+      async for event in ctx.agent.run_live(ctx):
+        yield event
+
+    async for event in self._exec_with_plugin(
+        invocation_context, session, execute
+    ):
       yield event
 
   def _find_agent_to_run(
@@ -355,9 +449,10 @@ class Runner:
     """Finds the agent to run to continue the session.
 
     A qualified agent must be either of:
+
     - The agent that returned a function call and the last user message is a
       function response to this function call.
-    - The root agent;
+    - The root agent.
     - An LlmAgent who replied last and is capable to transfer to any other agent
       in the agent hierarchy.
 
@@ -366,7 +461,8 @@ class Runner:
         root_agent: The root agent of the runner.
 
     Returns:
-      The agent of the last message in the session or the root agent.
+      The agent to run. (the active agent that should reply to the latest user
+      message)
     """
     # If the last event is a function response, should send this response to
     # the agent that returned the corressponding function call regardless the
@@ -395,8 +491,8 @@ class Runner:
   def _is_transferable_across_agent_tree(self, agent_to_run: BaseAgent) -> bool:
     """Whether the agent to run can transfer to any other agent in the agent tree.
 
-    This typically means all agent_to_run's parent through root agent can
-    transfer to their parent_agent.
+    This typically means all agent_to_run's ancestor can transfer to their
+    parent_agent all the way to the root_agent.
 
     Args:
         agent_to_run: The agent to check for transferability.
@@ -407,7 +503,7 @@ class Runner:
     agent = agent_to_run
     while agent:
       if not isinstance(agent, LlmAgent):
-        # Only LLM-based Agent can provider agent transfer capability.
+        # Only LLM-based Agent can provide agent transfer capability.
         return False
       if agent.disallow_transfer_to_parent:
         return False
@@ -450,6 +546,7 @@ class Runner:
         session_service=self.session_service,
         memory_service=self.memory_service,
         credential_service=self.credential_service,
+        plugin_manager=self.plugin_manager,
         invocation_id=invocation_id,
         agent=self.agent,
         session=session,
@@ -538,7 +635,13 @@ class InMemoryRunner(Runner):
         session service for the runner.
   """
 
-  def __init__(self, agent: BaseAgent, *, app_name: str = 'InMemoryRunner'):
+  def __init__(
+      self,
+      agent: BaseAgent,
+      *,
+      app_name: str = 'InMemoryRunner',
+      plugins: Optional[list[BasePlugin]] = None,
+  ):
     """Initializes the InMemoryRunner.
 
     Args:
@@ -551,6 +654,7 @@ class InMemoryRunner(Runner):
         app_name=app_name,
         agent=agent,
         artifact_service=InMemoryArtifactService(),
+        plugins=plugins,
         session_service=self._in_memory_session_service,
         memory_service=InMemoryMemoryService(),
     )

google/adk/sessions/database_session_service.py

@@ -137,7 +137,7 @@ class StorageSession(Base):
       DateTime(), default=func.now(), onupdate=func.now()
   )
 
-  storage_events: Mapped[list["StorageEvent"]] = relationship(
+  storage_events: Mapped[list[StorageEvent]] = relationship(
       "StorageEvent",
       back_populates="storage_session",
   )
@@ -373,11 +373,11 @@ class DatabaseSessionService(BaseSessionService):
     # 4. Build the session object with generated id
     # 5. Return the session
 
-    with self.database_session_factory() as session_factory:
+    with self.database_session_factory() as sql_session:
 
       # Fetch app and user states from storage
-      storage_app_state = session_factory.get(StorageAppState, (app_name))
-      storage_user_state = session_factory.get(
+      storage_app_state = sql_session.get(StorageAppState, (app_name))
+      storage_user_state = sql_session.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -387,12 +387,12 @@ class DatabaseSessionService(BaseSessionService):
       # Create state tables if not exist
       if not storage_app_state:
         storage_app_state = StorageAppState(app_name=app_name, state={})
-        session_factory.add(storage_app_state)
+        sql_session.add(storage_app_state)
       if not storage_user_state:
         storage_user_state = StorageUserState(
             app_name=app_name, user_id=user_id, state={}
         )
-        session_factory.add(storage_user_state)
+        sql_session.add(storage_user_state)
 
       # Extract state deltas
       app_state_delta, user_state_delta, session_state = _extract_state_delta(
@@ -416,10 +416,10 @@ class DatabaseSessionService(BaseSessionService):
           id=session_id,
           state=session_state,
       )
-      session_factory.add(storage_session)
-      session_factory.commit()
+      sql_session.add(storage_session)
+      sql_session.commit()
 
-      session_factory.refresh(storage_session)
+      sql_session.refresh(storage_session)
 
       # Merge states for response
       merged_state = _merge_state(app_state, user_state, session_state)
@@ -444,8 +444,8 @@ 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.database_session_factory() as session_factory:
-      storage_session = session_factory.get(
+    with self.database_session_factory() as sql_session:
+      storage_session = sql_session.get(
           StorageSession, (app_name, user_id, session_id)
       )
       if storage_session is None:
@@ -458,7 +458,7 @@ class DatabaseSessionService(BaseSessionService):
         timestamp_filter = True
 
       storage_events = (
-          session_factory.query(StorageEvent)
+          sql_session.query(StorageEvent)
           .filter(StorageEvent.app_name == app_name)
           .filter(StorageEvent.session_id == storage_session.id)
           .filter(StorageEvent.user_id == user_id)
@@ -473,8 +473,8 @@ class DatabaseSessionService(BaseSessionService):
       )
 
       # Fetch states from storage
-      storage_app_state = session_factory.get(StorageAppState, (app_name))
-      storage_user_state = session_factory.get(
+      storage_app_state = sql_session.get(StorageAppState, (app_name))
+      storage_user_state = sql_session.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -500,9 +500,9 @@ class DatabaseSessionService(BaseSessionService):
   async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
-    with self.database_session_factory() as session_factory:
+    with self.database_session_factory() as sql_session:
       results = (
-          session_factory.query(StorageSession)
+          sql_session.query(StorageSession)
           .filter(StorageSession.app_name == app_name)
           .filter(StorageSession.user_id == user_id)
           .all()
@@ -523,14 +523,14 @@ class DatabaseSessionService(BaseSessionService):
   async def delete_session(
       self, app_name: str, user_id: str, session_id: str
   ) -> None:
-    with self.database_session_factory() as session_factory:
+    with self.database_session_factory() as sql_session:
       stmt = delete(StorageSession).where(
           StorageSession.app_name == app_name,
           StorageSession.user_id == user_id,
           StorageSession.id == session_id,
       )
-      session_factory.execute(stmt)
-      session_factory.commit()
+      sql_session.execute(stmt)
+      sql_session.commit()
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
@@ -542,8 +542,8 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Check if timestamp is stale
     # 2. Update session attributes based on event config
     # 3. Store event to table
-    with self.database_session_factory() as session_factory:
-      storage_session = session_factory.get(
+    with self.database_session_factory() as sql_session:
+      storage_session = sql_session.get(
           StorageSession, (session.app_name, session.user_id, session.id)
       )
 
@@ -557,10 +557,8 @@ class DatabaseSessionService(BaseSessionService):
         )
 
       # Fetch states from storage
-      storage_app_state = session_factory.get(
-          StorageAppState, (session.app_name)
-      )
-      storage_user_state = session_factory.get(
+      storage_app_state = sql_session.get(StorageAppState, (session.app_name))
+      storage_user_state = sql_session.get(
           StorageUserState, (session.app_name, session.user_id)
       )
 
@@ -589,10 +587,10 @@ class DatabaseSessionService(BaseSessionService):
         session_state.update(session_state_delta)
         storage_session.state = session_state
 
-      session_factory.add(StorageEvent.from_event(session, event))
+      sql_session.add(StorageEvent.from_event(session, event))
 
-      session_factory.commit()
-      session_factory.refresh(storage_session)
+      sql_session.commit()
+      sql_session.refresh(storage_session)
 
       # Update timestamp with commit time
       session.last_update_time = storage_session.update_timestamp_tz

google/adk/sessions/vertex_ai_session_service.py

@@ -216,29 +216,36 @@ class VertexAiSessionService(BaseSessionService):
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events',
         request_dict={},
     )
-    list_events_api_response = _convert_api_response(list_events_api_response)
+    converted_api_response = _convert_api_response(list_events_api_response)
 
-    # Handles empty response case
-    if not list_events_api_response or list_events_api_response.get(
+    # Handles empty response case where there are no events to fetch
+    if not converted_api_response or converted_api_response.get(
         'httpHeaders', None
     ):
       return session
 
     session.events += [
         _from_api_event(event)
-        for event in list_events_api_response['sessionEvents']
+        for event in converted_api_response['sessionEvents']
     ]
 
-    while list_events_api_response.get('nextPageToken', None):
-      page_token = list_events_api_response.get('nextPageToken', None)
+    while converted_api_response.get('nextPageToken', None):
+      page_token = converted_api_response.get('nextPageToken', None)
       list_events_api_response = await api_client.async_request(
           http_method='GET',
           path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events?pageToken={page_token}',
           request_dict={},
       )
+      converted_api_response = _convert_api_response(list_events_api_response)
+
+      # Handles empty response case where there are no more events to fetch
+      if not converted_api_response or converted_api_response.get(
+          'httpHeaders', None
+      ):
+        break
       session.events += [
           _from_api_event(event)
-          for event in list_events_api_response['sessionEvents']
+          for event in converted_api_response['sessionEvents']
       ]
 
     session.events = [

google/adk/tools/agent_tool.py

@@ -116,6 +116,7 @@ class AgentTool(BaseTool):
         artifact_service=ForwardingArtifactService(tool_context),
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
+        credential_service=tool_context._invocation_context.credential_service,
     )
     session = await runner.session_service.create_session(
         app_name=self.agent.name,

google/adk/tools/apihub_tool/apihub_toolset.py

@@ -35,27 +35,25 @@ from .clients.apihub_client import APIHubClient
 class APIHubToolset(BaseToolset):
   """APIHubTool generates tools from a given API Hub resource.
 
-  Examples:
+  Examples::
 
-  ```
-  apihub_toolset = APIHubToolset(
-      apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
-      service_account_json="...",
-      tool_filter=lambda tool, ctx=None: tool.name in ('my_tool',
-      'my_other_tool')
-  )
-
-  # Get all available tools
-  agent = LlmAgent(tools=apihub_toolset)
+    apihub_toolset = APIHubToolset(
+        apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
+        service_account_json="...",
+        tool_filter=lambda tool, ctx=None: tool.name in ('my_tool',
+        'my_other_tool')
+    )
 
-  ```
+    # Get all available tools
+    agent = LlmAgent(tools=apihub_toolset)
 
   **apihub_resource_name** is the resource name from API Hub. It must include
-    API name, and can optionally include API version and spec name.
-    - If apihub_resource_name includes a spec resource name, the content of that
-      spec will be used for generating the tools.
-    - If apihub_resource_name includes only an api or a version name, the
-      first spec of the first version of that API will be used.
+  API name, and can optionally include API version and spec name.
+
+  - If apihub_resource_name includes a spec resource name, the content of that
+    spec will be used for generating the tools.
+  - If apihub_resource_name includes only an api or a version name, the
+    first spec of the first version of that API will be used.
   """
 
   def __init__(
@@ -78,44 +76,45 @@ class APIHubToolset(BaseToolset):
   ):
     """Initializes the APIHubTool with the given parameters.
 
-    Examples:
-    ```
-    apihub_toolset = APIHubToolset(
-        apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
-        service_account_json="...",
-    )
+    Examples::
 
-    # Get all available tools
-    agent = LlmAgent(tools=[apihub_toolset])
+      apihub_toolset = APIHubToolset(
+          apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
+          service_account_json="...",
+      )
 
-    apihub_toolset = APIHubToolset(
-        apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
-        service_account_json="...",
-        tool_filter = ['my_tool']
-    )
-    # Get a specific tool
-    agent = LlmAgent(tools=[
-        ...,
-        apihub_toolset,
-    ])
-    ```
+      # Get all available tools
+      agent = LlmAgent(tools=[apihub_toolset])
+
+      apihub_toolset = APIHubToolset(
+          apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
+          service_account_json="...",
+          tool_filter = ['my_tool']
+      )
+      # Get a specific tool
+      agent = LlmAgent(tools=[
+          ...,
+          apihub_toolset,
+      ])
 
     **apihub_resource_name** is the resource name from API Hub. It must include
     API name, and can optionally include API version and spec name.
+
     - If apihub_resource_name includes a spec resource name, the content of that
       spec will be used for generating the tools.
     - If apihub_resource_name includes only an api or a version name, the
       first spec of the first version of that API will be used.
 
     Example:
+
     * projects/xxx/locations/us-central1/apis/apiname/...
     * https://console.cloud.google.com/apigee/api-hub/apis/apiname?project=xxx
 
     Args:
         apihub_resource_name: The resource name of the API in API Hub.
-          Example: `projects/test-project/locations/us-central1/apis/test-api`.
-        access_token: Google Access token. Generate with gcloud cli `gcloud auth
-          auth print-access-token`. Used for fetching API Specs from API Hub.
+          Example: ``projects/test-project/locations/us-central1/apis/test-api``.
+        access_token: Google Access token. Generate with gcloud cli
+          ``gcloud auth auth print-access-token``. Used for fetching API Specs from API Hub.
         service_account_json: The service account config as a json string.
           Required if not using default service credential. It is used for
           creating the API Hub client and fetching the API Specs from API Hub.