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

google/adk/agents/active_streaming_tool.py

@@ -30,6 +30,7 @@ class ActiveStreamingTool(BaseModel):
       arbitrary_types_allowed=True,
       extra='forbid',
   )
+  """The pydantic model config."""
 
   task: Optional[asyncio.Task] = None
   """The active task of this streaming tool."""

google/adk/agents/base_agent.py

@@ -14,12 +14,15 @@
 
 from __future__ import annotations
 
+import inspect
 from typing import Any
 from typing import AsyncGenerator
+from typing import Awaitable
 from typing import Callable
 from typing import final
 from typing import Optional
 from typing import TYPE_CHECKING
+from typing import Union
 
 from google.genai import types
 from opentelemetry import trace
@@ -28,6 +31,7 @@ from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import field_validator
 from typing_extensions import override
+from typing_extensions import TypeAlias
 
 from ..events.event import Event
 from .callback_context import CallbackContext
@@ -37,27 +41,20 @@ if TYPE_CHECKING:
 
 tracer = trace.get_tracer('gcp.vertex.agent')
 
-BeforeAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
-"""Callback signature that is invoked before the agent run.
+_SingleAgentCallback: TypeAlias = Callable[
+    [CallbackContext],
+    Union[Awaitable[Optional[types.Content]], Optional[types.Content]],
+]
 
-Args:
-  callback_context: MUST be named 'callback_context' (enforced).
+BeforeAgentCallback: TypeAlias = Union[
+    _SingleAgentCallback,
+    list[_SingleAgentCallback],
+]
 
-Returns:
-  The content to return to the user. When set, the agent run will be skipped and
-  the provided content will be returned to user.
-"""
-
-AfterAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
-"""Callback signature that is invoked after the agent run.
-
-Args:
-  callback_context: MUST be named 'callback_context' (enforced).
-
-Returns:
-  The content to return to the user. When set, the provided content will be
-  appended to event history as agent response.
-"""
+AfterAgentCallback: TypeAlias = Union[
+    _SingleAgentCallback,
+    list[_SingleAgentCallback],
+]
 
 
 class BaseAgent(BaseModel):
@@ -67,6 +64,7 @@ class BaseAgent(BaseModel):
       arbitrary_types_allowed=True,
       extra='forbid',
   )
+  """The pydantic model config."""
 
   name: str
   """The agent's name.
@@ -95,24 +93,32 @@ class BaseAgent(BaseModel):
   """The sub-agents of this agent."""
 
   before_agent_callback: Optional[BeforeAgentCallback] = None
-  """Callback signature that is invoked before the agent run.
+  """Callback or list of callbacks to be invoked before the agent run.
+
+  When a list of callbacks is provided, the callbacks will be called in the
+  order they are listed until a callback does not return None.
 
   Args:
     callback_context: MUST be named 'callback_context' (enforced).
 
   Returns:
-    The content to return to the user. When set, the agent run will be skipped
-    and the provided content will be returned to user.
+    Optional[types.Content]: The content to return to the user.
+      When the content is present, the agent run will be skipped and the
+      provided content will be returned to user.
   """
   after_agent_callback: Optional[AfterAgentCallback] = None
-  """Callback signature that is invoked after the agent run.
+  """Callback or list of callbacks to be invoked after the agent run.
+
+  When a list of callbacks is provided, the callbacks will be called in the
+  order they are listed until a callback does not return None.
 
   Args:
     callback_context: MUST be named 'callback_context' (enforced).
 
   Returns:
-    The content to return to the user. When set, the provided content will be
-    appended to event history as agent response.
+    Optional[types.Content]: The content to return to the user.
+      When the content is present, the provided content will be used as agent
+      response and appended to event history as agent response.
   """
 
   @final
@@ -133,7 +139,7 @@ class BaseAgent(BaseModel):
     with tracer.start_as_current_span(f'agent_run [{self.name}]'):
       ctx = self._create_invocation_context(parent_context)
 
-      if event := self.__handle_before_agent_callback(ctx):
+      if event := await self.__handle_before_agent_callback(ctx):
         yield event
       if ctx.end_invocation:
         return
@@ -144,7 +150,7 @@ class BaseAgent(BaseModel):
       if ctx.end_invocation:
         return
 
-      if event := self.__handle_after_agent_callback(ctx):
+      if event := await self.__handle_after_agent_callback(ctx):
         yield event
 
   @final
@@ -244,7 +250,31 @@ class BaseAgent(BaseModel):
       invocation_context.branch = f'{parent_context.branch}.{self.name}'
     return invocation_context
 
-  def __handle_before_agent_callback(
+  @property
+  def canonical_before_agent_callbacks(self) -> list[_SingleAgentCallback]:
+    """The resolved self.before_agent_callback field as a list of _SingleAgentCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.before_agent_callback:
+      return []
+    if isinstance(self.before_agent_callback, list):
+      return self.before_agent_callback
+    return [self.before_agent_callback]
+
+  @property
+  def canonical_after_agent_callbacks(self) -> list[_SingleAgentCallback]:
+    """The resolved self.after_agent_callback field as a list of _SingleAgentCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.after_agent_callback:
+      return []
+    if isinstance(self.after_agent_callback, list):
+      return self.after_agent_callback
+    return [self.after_agent_callback]
+
+  async def __handle_before_agent_callback(
       self, ctx: InvocationContext
   ) -> Optional[Event]:
     """Runs the before_agent_callback if it exists.
@@ -254,24 +284,27 @@ class BaseAgent(BaseModel):
     """
     ret_event = None
 
-    if not isinstance(self.before_agent_callback, Callable):
+    if not self.canonical_before_agent_callbacks:
       return ret_event
 
     callback_context = CallbackContext(ctx)
-    before_agent_callback_content = self.before_agent_callback(
-        callback_context=callback_context
-    )
 
-    if before_agent_callback_content:
-      ret_event = Event(
-          invocation_id=ctx.invocation_id,
-          author=self.name,
-          branch=ctx.branch,
-          content=before_agent_callback_content,
-          actions=callback_context._event_actions,
+    for callback in self.canonical_before_agent_callbacks:
+      before_agent_callback_content = callback(
+          callback_context=callback_context
       )
-      ctx.end_invocation = True
-      return ret_event
+      if inspect.isawaitable(before_agent_callback_content):
+        before_agent_callback_content = await before_agent_callback_content
+      if before_agent_callback_content:
+        ret_event = Event(
+            invocation_id=ctx.invocation_id,
+            author=self.name,
+            branch=ctx.branch,
+            content=before_agent_callback_content,
+            actions=callback_context._event_actions,
+        )
+        ctx.end_invocation = True
+        return ret_event
 
     if callback_context.state.has_delta():
       ret_event = Event(
@@ -283,7 +316,7 @@ class BaseAgent(BaseModel):
 
     return ret_event
 
-  def __handle_after_agent_callback(
+  async def __handle_after_agent_callback(
       self, invocation_context: InvocationContext
   ) -> Optional[Event]:
     """Runs the after_agent_callback if it exists.
@@ -293,15 +326,26 @@ class BaseAgent(BaseModel):
     """
     ret_event = None
 
-    if not isinstance(self.after_agent_callback, Callable):
+    if not self.canonical_after_agent_callbacks:
       return ret_event
 
     callback_context = CallbackContext(invocation_context)
-    after_agent_callback_content = self.after_agent_callback(
-        callback_context=callback_context
-    )
 
-    if after_agent_callback_content or callback_context.state.has_delta():
+    for callback in self.canonical_after_agent_callbacks:
+      after_agent_callback_content = callback(callback_context=callback_context)
+      if inspect.isawaitable(after_agent_callback_content):
+        after_agent_callback_content = await after_agent_callback_content
+      if after_agent_callback_content:
+        ret_event = Event(
+            invocation_id=invocation_context.invocation_id,
+            author=self.name,
+            branch=invocation_context.branch,
+            content=after_agent_callback_content,
+            actions=callback_context._event_actions,
+        )
+        return ret_event
+
+    if callback_context.state.has_delta():
       ret_event = Event(
           invocation_id=invocation_context.invocation_id,
           author=self.name,

google/adk/agents/base_agent.py.orig

@@ -0,0 +1,330 @@
+# 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
+
+from typing import Any
+from typing import AsyncGenerator
+from typing import Callable
+from typing import final
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from google.genai import types
+from opentelemetry import trace
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import field_validator
+from typing_extensions import override
+
+from ..events.event import Event
+from .callback_context import CallbackContext
+
+if TYPE_CHECKING:
+  from .invocation_context import InvocationContext
+
+tracer = trace.get_tracer('gcp.vertex.agent')
+
+BeforeAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
+
+
+AfterAgentCallback = Callable[[CallbackContext], Optional[types.Content]]
+
+
+class BaseAgent(BaseModel):
+  """Base class for all agents in Agent Development Kit."""
+
+  model_config = ConfigDict(
+      arbitrary_types_allowed=True,
+      extra='forbid',
+  )
+
+  name: str
+  """The agent's name.
+
+  Agent name must be a Python identifier and unique within the agent tree.
+  Agent name cannot be "user", since it's reserved for end-user's input.
+  """
+
+  description: str = ''
+  """Description about the agent's capability.
+
+  The model uses this to determine whether to delegate control to the agent.
+  One-line description is enough and preferred.
+  """
+
+  parent_agent: Optional[BaseAgent] = Field(default=None, init=False)
+  """The parent agent of this agent.
+
+  Note that an agent can ONLY be added as sub-agent once.
+
+  If you want to add one agent twice as sub-agent, consider to create two agent
+  instances with identical config, but with different name and add them to the
+  agent tree.
+  """
+  sub_agents: list[BaseAgent] = Field(default_factory=list)
+  """The sub-agents of this agent."""
+
+  before_agent_callback: Optional[BeforeAgentCallback] = None
+  """Callback signature that is invoked before the agent run.
+
+  Args:
+    callback_context: MUST be named 'callback_context' (enforced).
+
+  Returns:
+    Optional[types.Content]: The content to return to the user.
+      When the content is present, the agent run will be skipped and the
+      provided content will be returned to user.
+  """
+  after_agent_callback: Optional[AfterAgentCallback] = None
+  """Callback signature that is invoked after the agent run.
+
+  Args:
+    callback_context: MUST be named 'callback_context' (enforced).
+
+  Returns:
+    Optional[types.Content]: The content to return to the user.
+      When the content is present, the provided content will be used as agent
+      response and appended to event history as agent response.
+  """
+
+  @final
+  async def run_async(
+      self,
+      parent_context: InvocationContext,
+  ) -> AsyncGenerator[Event, None]:
+    """Entry method to run an agent via text-based conversation.
+
+    Args:
+      parent_context: InvocationContext, the invocation context of the parent
+        agent.
+
+    Yields:
+      Event: the events generated by the agent.
+    """
+
+    with tracer.start_as_current_span(f'agent_run [{self.name}]'):
+      ctx = self._create_invocation_context(parent_context)
+
+      if event := self.__handle_before_agent_callback(ctx):
+        yield event
+      if ctx.end_invocation:
+        return
+
+      async for event in self._run_async_impl(ctx):
+        yield event
+
+      if ctx.end_invocation:
+        return
+
+      if event := self.__handle_after_agent_callback(ctx):
+        yield event
+
+  @final
+  async def run_live(
+      self,
+      parent_context: InvocationContext,
+  ) -> AsyncGenerator[Event, None]:
+    """Entry method to run an agent via video/audio-based conversation.
+
+    Args:
+      parent_context: InvocationContext, the invocation context of the parent
+        agent.
+
+    Yields:
+      Event: the events generated by the agent.
+    """
+    with tracer.start_as_current_span(f'agent_run [{self.name}]'):
+      ctx = self._create_invocation_context(parent_context)
+      # TODO(hangfei): support before/after_agent_callback
+
+      async for event in self._run_live_impl(ctx):
+        yield event
+
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    """Core logic to run this agent via text-based conversation.
+
+    Args:
+      ctx: InvocationContext, the invocation context for this agent.
+
+    Yields:
+      Event: the events generated by the agent.
+    """
+    raise NotImplementedError(
+        f'_run_async_impl for {type(self)} is not implemented.'
+    )
+    yield  # AsyncGenerator requires having at least one yield statement
+
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    """Core logic to run this agent via video/audio-based conversation.
+
+    Args:
+      ctx: InvocationContext, the invocation context for this agent.
+
+    Yields:
+      Event: the events generated by the agent.
+    """
+    raise NotImplementedError(
+        f'_run_live_impl for {type(self)} is not implemented.'
+    )
+    yield  # AsyncGenerator requires having at least one yield statement
+
+  @property
+  def root_agent(self) -> BaseAgent:
+    """Gets the root agent of this agent."""
+    root_agent = self
+    while root_agent.parent_agent is not None:
+      root_agent = root_agent.parent_agent
+    return root_agent
+
+  def find_agent(self, name: str) -> Optional[BaseAgent]:
+    """Finds the agent with the given name in this agent and its descendants.
+
+    Args:
+      name: The name of the agent to find.
+
+    Returns:
+      The agent with the matching name, or None if no such agent is found.
+    """
+    if self.name == name:
+      return self
+    return self.find_sub_agent(name)
+
+  def find_sub_agent(self, name: str) -> Optional[BaseAgent]:
+    """Finds the agent with the given name in this agent's descendants.
+
+    Args:
+      name: The name of the agent to find.
+
+    Returns:
+      The agent with the matching name, or None if no such agent is found.
+    """
+    for sub_agent in self.sub_agents:
+      if result := sub_agent.find_agent(name):
+        return result
+    return None
+
+  def _create_invocation_context(
+      self, parent_context: InvocationContext
+  ) -> InvocationContext:
+    """Creates a new invocation context for this agent."""
+    invocation_context = parent_context.model_copy(update={'agent': self})
+    if parent_context.branch:
+      invocation_context.branch = f'{parent_context.branch}.{self.name}'
+    return invocation_context
+
+  def __handle_before_agent_callback(
+      self, ctx: InvocationContext
+  ) -> Optional[Event]:
+    """Runs the before_agent_callback if it exists.
+
+    Returns:
+      Optional[Event]: an event if callback provides content or changed state.
+    """
+    ret_event = None
+
+    if not isinstance(self.before_agent_callback, Callable):
+      return ret_event
+
+    callback_context = CallbackContext(ctx)
+    before_agent_callback_content = self.before_agent_callback(
+        callback_context=callback_context
+    )
+
+    if before_agent_callback_content:
+      ret_event = Event(
+          invocation_id=ctx.invocation_id,
+          author=self.name,
+          branch=ctx.branch,
+          content=before_agent_callback_content,
+          actions=callback_context._event_actions,
+      )
+      ctx.end_invocation = True
+      return ret_event
+
+    if callback_context.state.has_delta():
+      ret_event = Event(
+          invocation_id=ctx.invocation_id,
+          author=self.name,
+          branch=ctx.branch,
+          actions=callback_context._event_actions,
+      )
+
+    return ret_event
+
+  def __handle_after_agent_callback(
+      self, invocation_context: InvocationContext
+  ) -> Optional[Event]:
+    """Runs the after_agent_callback if it exists.
+
+    Returns:
+      Optional[Event]: an event if callback provides content or changed state.
+    """
+    ret_event = None
+
+    if not isinstance(self.after_agent_callback, Callable):
+      return ret_event
+
+    callback_context = CallbackContext(invocation_context)
+    after_agent_callback_content = self.after_agent_callback(
+        callback_context=callback_context
+    )
+
+    if after_agent_callback_content or callback_context.state.has_delta():
+      ret_event = Event(
+          invocation_id=invocation_context.invocation_id,
+          author=self.name,
+          branch=invocation_context.branch,
+          content=after_agent_callback_content,
+          actions=callback_context._event_actions,
+      )
+
+    return ret_event
+
+  @override
+  def model_post_init(self, __context: Any) -> None:
+    self.__set_parent_agent_for_sub_agents()
+
+  @field_validator('name', mode='after')
+  @classmethod
+  def __validate_name(cls, value: str):
+    if not value.isidentifier():
+      raise ValueError(
+          f'Found invalid agent name: `{value}`.'
+          ' Agent name must be a valid identifier. It should start with a'
+          ' letter (a-z, A-Z) or an underscore (_), and can only contain'
+          ' letters, digits (0-9), and underscores.'
+      )
+    if value == 'user':
+      raise ValueError(
+          "Agent name cannot be `user`. `user` is reserved for end-user's"
+          ' input.'
+      )
+    return value
+
+  def __set_parent_agent_for_sub_agents(self) -> BaseAgent:
+    for sub_agent in self.sub_agents:
+      if sub_agent.parent_agent is not None:
+        raise ValueError(
+            f'Agent `{sub_agent.name}` already has a parent agent, current'
+            f' parent: `{sub_agent.parent_agent.name}`, trying to add:'
+            f' `{self.name}`'
+        )
+      sub_agent.parent_agent = self
+    return self

google/adk/agents/callback_context.py

@@ -60,12 +60,7 @@ class CallbackContext(ReadonlyContext):
     """
     return self._state
 
-  @property
-  def user_content(self) -> Optional[types.Content]:
-    """The user content that started this invocation. READONLY field."""
-    return self._invocation_context.user_content
-
-  def load_artifact(
+  async def load_artifact(
       self, filename: str, version: Optional[int] = None
   ) -> Optional[types.Part]:
     """Loads an artifact attached to the current session.
@@ -80,7 +75,7 @@ class CallbackContext(ReadonlyContext):
     """
     if self._invocation_context.artifact_service is None:
       raise ValueError("Artifact service is not initialized.")
-    return self._invocation_context.artifact_service.load_artifact(
+    return await self._invocation_context.artifact_service.load_artifact(
         app_name=self._invocation_context.app_name,
         user_id=self._invocation_context.user_id,
         session_id=self._invocation_context.session.id,
@@ -88,7 +83,7 @@ class CallbackContext(ReadonlyContext):
         version=version,
     )
 
-  def save_artifact(self, filename: str, artifact: types.Part) -> int:
+  async def save_artifact(self, filename: str, artifact: types.Part) -> int:
     """Saves an artifact and records it as delta for the current session.
 
     Args:
@@ -100,7 +95,7 @@ class CallbackContext(ReadonlyContext):
     """
     if self._invocation_context.artifact_service is None:
       raise ValueError("Artifact service is not initialized.")
-    version = self._invocation_context.artifact_service.save_artifact(
+    version = await self._invocation_context.artifact_service.save_artifact(
         app_name=self._invocation_context.app_name,
         user_id=self._invocation_context.user_id,
         session_id=self._invocation_context.session.id,

google/adk/agents/invocation_context.py

@@ -110,6 +110,7 @@ class InvocationContext(BaseModel):
       arbitrary_types_allowed=True,
       extra="forbid",
   )
+  """The pydantic model config."""
 
   artifact_service: Optional[BaseArtifactService] = None
   session_service: BaseSessionService

google/adk/agents/langgraph_agent.py

@@ -53,6 +53,7 @@ class LangGraphAgent(BaseAgent):
   model_config = ConfigDict(
       arbitrary_types_allowed=True,
   )
+  """The pydantic model config."""
 
   graph: CompiledGraph
 

google/adk/agents/live_request_queue.py

@@ -24,6 +24,7 @@ class LiveRequest(BaseModel):
   """Request send to live agents."""
 
   model_config = ConfigDict(ser_json_bytes='base64', val_json_bytes='base64')
+  """The pydantic model config."""
 
   content: Optional[types.Content] = None
   """If set, send the content to the model in turn-by-turn mode."""