Agent Development Kit(ADK)

An easy-to-use and powerful framework to build AI agents.

src/google/adk/runners.py

@@ -0,0 +1,456 @@
+# 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 __future__ import annotations
+
+import asyncio
+import logging
+import queue
+import threading
+from typing import AsyncGenerator
+from typing import Generator
+from typing import Optional
+
+from deprecated import deprecated
+from google.genai import types
+
+from .agents.active_streaming_tool import ActiveStreamingTool
+from .agents.base_agent import BaseAgent
+from .agents.invocation_context import InvocationContext
+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 .events.event import Event
+from .memory.base_memory_service import BaseMemoryService
+from .memory.in_memory_memory_service import InMemoryMemoryService
+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
+
+logger = logging.getLogger(__name__)
+
+
+class Runner:
+  """The Runner class is used to run agents.
+
+  It manages the execution of an agent within a session, handling message
+  processing, event generation, and interaction with various services like
+  artifact storage, session management, and memory.
+
+  Attributes:
+      app_name: The application name of the runner.
+      agent: The root agent to run.
+      artifact_service: The artifact service for the runner.
+      session_service: The session service for the runner.
+      memory_service: The memory service for the runner.
+  """
+
+  app_name: str
+  """The app name of the runner."""
+  agent: BaseAgent
+  """The root agent to run."""
+  artifact_service: Optional[BaseArtifactService] = None
+  """The artifact service for the runner."""
+  session_service: BaseSessionService
+  """The session service for the runner."""
+  memory_service: Optional[BaseMemoryService] = None
+  """The memory service for the runner."""
+
+  def __init__(
+      self,
+      *,
+      app_name: str,
+      agent: BaseAgent,
+      artifact_service: Optional[BaseArtifactService] = None,
+      session_service: BaseSessionService,
+      memory_service: Optional[BaseMemoryService] = None,
+  ):
+    """Initializes the Runner.
+
+    Args:
+        app_name: The application name of the runner.
+        agent: The root agent to run.
+        artifact_service: The artifact service for the runner.
+        session_service: The session service for the runner.
+        memory_service: The memory service for the runner.
+    """
+    self.app_name = app_name
+    self.agent = agent
+    self.artifact_service = artifact_service
+    self.session_service = session_service
+    self.memory_service = memory_service
+
+  def run(
+      self,
+      *,
+      user_id: str,
+      session_id: str,
+      new_message: types.Content,
+      run_config: RunConfig = RunConfig(),
+  ) -> Generator[Event, None, None]:
+    """Runs the agent.
+
+    NOTE: This sync interface is only for local testing and convenience purpose.
+    Consider to use `run_async` for production usage.
+
+    Args:
+      user_id: The user ID of the session.
+      session_id: The session ID of the session.
+      new_message: A new message to append to the session.
+      run_config: The run config for the agent.
+
+    Yields:
+      The events generated by the agent.
+    """
+    event_queue = queue.Queue()
+
+    async def _invoke_run_async():
+      try:
+        async for event in self.run_async(
+            user_id=user_id,
+            session_id=session_id,
+            new_message=new_message,
+            run_config=run_config,
+        ):
+          event_queue.put(event)
+      finally:
+        event_queue.put(None)
+
+    def _asyncio_thread_main():
+      try:
+        asyncio.run(_invoke_run_async())
+      finally:
+        event_queue.put(None)
+
+    thread = threading.Thread(target=_asyncio_thread_main)
+    thread.start()
+
+    # consumes and re-yield the events from background thread.
+    while True:
+      event = event_queue.get()
+      if event is None:
+        break
+      else:
+        yield event
+
+    thread.join()
+
+  async def run_async(
+      self,
+      *,
+      user_id: str,
+      session_id: str,
+      new_message: types.Content,
+      run_config: RunConfig = RunConfig(),
+  ) -> AsyncGenerator[Event, None]:
+    """Main entry method to run the agent in this runner.
+
+    Args:
+      user_id: The user ID of the session.
+      session_id: The session ID of the session.
+      new_message: A new message to append to the session.
+      run_config: The run config for the agent.
+
+    Yields:
+      The events generated by the agent.
+    """
+    with tracer.start_as_current_span('invocation'):
+      session = 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(
+          session,
+          new_message=new_message,
+          run_config=run_config,
+      )
+      root_agent = self.agent
+
+      if new_message:
+        self._append_new_message_to_session(
+            session,
+            new_message,
+            invocation_context,
+            run_config.save_input_blobs_as_artifacts,
+        )
+
+      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)
+        yield event
+
+  def _append_new_message_to_session(
+      self,
+      session: Session,
+      new_message: types.Content,
+      invocation_context: InvocationContext,
+      save_input_blobs_as_artifacts: bool = False,
+  ):
+    """Appends a new message to the session.
+
+    Args:
+        session: The session to append the message to.
+        new_message: The new message to append.
+        invocation_context: The invocation context for the message.
+        save_input_blobs_as_artifacts: Whether to save input blobs as artifacts.
+    """
+    if not new_message.parts:
+      raise ValueError('No parts in the new_message.')
+
+    if self.artifact_service and save_input_blobs_as_artifacts:
+      # The runner directly saves the artifacts (if applicable) in the
+      # user message and replaces the artifact data with a file name
+      # placeholder.
+      for i, part in enumerate(new_message.parts):
+        if part.inline_data is None:
+          continue
+        file_name = f'artifact_{invocation_context.invocation_id}_{i}'
+        self.artifact_service.save_artifact(
+            app_name=self.app_name,
+            user_id=session.user_id,
+            session_id=session.id,
+            filename=file_name,
+            artifact=part,
+        )
+        new_message.parts[i] = types.Part(
+            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,
+    )
+    self.session_service.append_event(session=session, event=event)
+
+  async def run_live(
+      self,
+      *,
+      session: Session,
+      live_request_queue: LiveRequestQueue,
+      run_config: RunConfig = RunConfig(),
+  ) -> AsyncGenerator[Event, None]:
+    """Runs the agent in live mode (experimental feature).
+
+    Args:
+        session: The session to use.
+        live_request_queue: The queue for live requests.
+        run_config: The run config for the agent.
+
+    Yields:
+        The events generated by the agent.
+    """
+    # TODO: right now, only works for a single audio agent without FC.
+    invocation_context = self._new_invocation_context_for_live(
+        session,
+        live_request_queue=live_request_queue,
+        run_config=run_config,
+    )
+
+    root_agent = self.agent
+    invocation_context.agent = self._find_agent_to_run(session, root_agent)
+
+    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
+          )
+
+    async for event in invocation_context.agent.run_live(invocation_context):
+      self.session_service.append_event(session=session, event=event)
+      yield event
+
+  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:
+      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:
+    """Finds the agent to run to continue the session.
+
+    A qualified agent must be either of:
+    - The root agent;
+    - An LlmAgent who replied last and is capable to transfer to any other agent
+      in the agent hierarchy.
+
+    Args:
+        session: The session to find the agent for.
+        root_agent: The root agent of the runner.
+
+    Returns:
+      The agent of the last message in the session or the root agent.
+    """
+    for event in filter(lambda e: e.author != 'user', reversed(session.events)):
+      if event.author == root_agent.name:
+        # Found root agent.
+        return root_agent
+      if not (agent := root_agent.find_sub_agent(event.author)):
+        # Agent not found, continue looking.
+        logger.warning(
+            'Event from an unknown agent: %s, event id: %s',
+            event.author,
+            event.id,
+        )
+        continue
+      if self._is_transferable_across_agent_tree(agent):
+        return agent
+    # Falls back to root agent if no suitable agents are found in the session.
+    return root_agent
+
+  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.
+
+    Args:
+        agent_to_run: The agent to check for transferability.
+
+    Returns:
+        True if the agent can transfer, False otherwise.
+    """
+    agent = agent_to_run
+    while agent:
+      if not isinstance(agent, LlmAgent):
+        # Only LLM-based Agent can provider agent transfer capability.
+        return False
+      if agent.disallow_transfer_to_parent:
+        return False
+      agent = agent.parent_agent
+    return True
+
+  def _new_invocation_context(
+      self,
+      session: Session,
+      *,
+      new_message: Optional[types.Content] = None,
+      live_request_queue: Optional[LiveRequestQueue] = None,
+      run_config: RunConfig = RunConfig(),
+  ) -> InvocationContext:
+    """Creates a new invocation context.
+
+    Args:
+        session: The session for the context.
+        new_message: The new message for the context.
+        live_request_queue: The live request queue for the context.
+        run_config: The run config for the context.
+
+    Returns:
+        The new invocation context.
+    """
+    invocation_id = new_invocation_context_id()
+
+    if run_config.support_cfc and isinstance(self.agent, LlmAgent):
+      model_name = self.agent.canonical_model.model
+      if not model_name.startswith('gemini-2'):
+        raise ValueError(
+            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)
+
+    return InvocationContext(
+        artifact_service=self.artifact_service,
+        session_service=self.session_service,
+        memory_service=self.memory_service,
+        invocation_id=invocation_id,
+        agent=self.agent,
+        session=session,
+        user_content=new_message,
+        live_request_queue=live_request_queue,
+        run_config=run_config,
+    )
+
+  def _new_invocation_context_for_live(
+      self,
+      session: Session,
+      *,
+      live_request_queue: Optional[LiveRequestQueue] = None,
+      run_config: RunConfig = RunConfig(),
+  ) -> InvocationContext:
+    """Creates a new invocation context for live multi-agent."""
+
+    # For live multi-agent, we need model's text transcription as context for
+    # next agent.
+    if self.agent.sub_agents and live_request_queue:
+      if not run_config.response_modalities:
+        # default
+        run_config.response_modalities = ['AUDIO', 'TEXT']
+      elif 'TEXT' not in run_config.response_modalities:
+        run_config.response_modalities.append('TEXT')
+    return self._new_invocation_context(
+        session,
+        live_request_queue=live_request_queue,
+        run_config=run_config,
+    )
+
+
+class InMemoryRunner(Runner):
+  """An in-memory Runner for testing and development.
+
+  This runner uses in-memory implementations for artifact, session, and memory
+  services, providing a lightweight and self-contained environment for agent
+  execution.
+
+  Attributes:
+      agent: The root agent to run.
+      app_name: The application name of the runner. Defaults to
+        'InMemoryRunner'.
+  """
+
+  def __init__(self, agent: LlmAgent, *, app_name: str = 'InMemoryRunner'):
+    """Initializes the InMemoryRunner.
+
+    Args:
+        agent: The root agent to run.
+        app_name: The application name of the runner. Defaults to
+          'InMemoryRunner'.
+    """
+    super().__init__(
+        app_name=app_name,
+        agent=agent,
+        artifact_service=InMemoryArtifactService(),
+        session_service=InMemorySessionService(),
+        memory_service=InMemoryMemoryService(),
+    )