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

google/adk/agents/llm_agent.py

@@ -14,8 +14,15 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
-from typing import Any, AsyncGenerator, Awaitable, Callable, Literal, Optional, Union
+from typing import Any
+from typing import AsyncGenerator
+from typing import Awaitable
+from typing import Callable
+from typing import Literal
+from typing import Optional
+from typing import Union
 
 from google.genai import types
 from pydantic import BaseModel
@@ -38,6 +45,7 @@ from ..models.llm_response import LlmResponse
 from ..models.registry import LLMRegistry
 from ..planners.base_planner import BasePlanner
 from ..tools.base_tool import BaseTool
+from ..tools.base_toolset import BaseToolset
 from ..tools.function_tool import FunctionTool
 from ..tools.tool_context import ToolContext
 from .base_agent import BaseAgent
@@ -45,39 +53,65 @@ from .callback_context import CallbackContext
 from .invocation_context import InvocationContext
 from .readonly_context import ReadonlyContext
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
+_SingleBeforeModelCallback: TypeAlias = Callable[
+    [CallbackContext, LlmRequest],
+    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
+]
 
-BeforeModelCallback: TypeAlias = Callable[
-    [CallbackContext, LlmRequest], Optional[LlmResponse]
+BeforeModelCallback: TypeAlias = Union[
+    _SingleBeforeModelCallback,
+    list[_SingleBeforeModelCallback],
 ]
-AfterModelCallback: TypeAlias = Callable[
+
+_SingleAfterModelCallback: TypeAlias = Callable[
     [CallbackContext, LlmResponse],
-    Optional[LlmResponse],
+    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
+]
+
+AfterModelCallback: TypeAlias = Union[
+    _SingleAfterModelCallback,
+    list[_SingleAfterModelCallback],
 ]
-BeforeToolCallback: TypeAlias = Callable[
+
+_SingleBeforeToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext],
     Union[Awaitable[Optional[dict]], Optional[dict]],
 ]
-AfterToolCallback: TypeAlias = Callable[
+
+BeforeToolCallback: TypeAlias = Union[
+    _SingleBeforeToolCallback,
+    list[_SingleBeforeToolCallback],
+]
+
+_SingleAfterToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext, dict],
     Union[Awaitable[Optional[dict]], Optional[dict]],
 ]
 
-InstructionProvider: TypeAlias = Callable[[ReadonlyContext], str]
+AfterToolCallback: TypeAlias = Union[
+    _SingleAfterToolCallback,
+    list[_SingleAfterToolCallback],
+]
+
+InstructionProvider: TypeAlias = Callable[
+    [ReadonlyContext], Union[str, Awaitable[str]]
+]
 
-ToolUnion: TypeAlias = Union[Callable, BaseTool]
+ToolUnion: TypeAlias = Union[Callable, BaseTool, BaseToolset]
 ExamplesUnion = Union[list[Example], BaseExampleProvider]
 
 
-def _convert_tool_union_to_tool(
-    tool_union: ToolUnion,
-) -> BaseTool:
-  return (
-      tool_union
-      if isinstance(tool_union, BaseTool)
-      else FunctionTool(tool_union)
-  )
+async def _convert_tool_union_to_tools(
+    tool_union: ToolUnion, ctx: ReadonlyContext
+) -> list[BaseTool]:
+  if isinstance(tool_union, BaseTool):
+    return [tool_union]
+  if isinstance(tool_union, Callable):
+    return [FunctionTool(func=tool_union)]
+
+  return await tool_union.get_tools(ctx)
 
 
 class LlmAgent(BaseAgent):
@@ -116,7 +150,12 @@ class LlmAgent(BaseAgent):
 
   # LLM-based agent transfer configs - Start
   disallow_transfer_to_parent: bool = False
-  """Disallows LLM-controlled transferring to the parent agent."""
+  """Disallows LLM-controlled transferring to the parent agent.
+
+  NOTE: Setting this as True also prevents this agent to continue reply to the
+  end-user. This behavior prevents one-way transfer, in which end-user may be
+  stuck with one agent that cannot transfer to other agents in the agent tree.
+  """
   disallow_transfer_to_peers: bool = False
   """Disallows LLM-controlled transferring to the peer agents."""
   # LLM-based agent transfer configs - End
@@ -161,8 +200,7 @@ class LlmAgent(BaseAgent):
 
   Check out available code executions in `google.adk.code_executor` package.
 
-  NOTE: to use model's built-in code executor, don't set this field, add
-  `google.adk.tools.built_in_code_execution` to tools instead.
+  NOTE: to use model's built-in code executor, use the `BuiltInCodeExecutor`.
   """
   # Advance features - End
 
@@ -173,7 +211,11 @@ class LlmAgent(BaseAgent):
 
   # Callbacks - Start
   before_model_callback: Optional[BeforeModelCallback] = None
-  """Called before calling the LLM.
+  """Callback or list of callbacks to be called before calling the LLM.
+
+  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: CallbackContext,
     llm_request: LlmRequest, The raw model request. Callback can mutate the
@@ -184,7 +226,10 @@ class LlmAgent(BaseAgent):
     skipped and the provided content will be returned to user.
   """
   after_model_callback: Optional[AfterModelCallback] = None
-  """Called after calling LLM.
+  """Callback or list of callbacks to be called after calling the LLM.
+
+  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: CallbackContext,
@@ -195,7 +240,10 @@ class LlmAgent(BaseAgent):
     will be ignored and the provided content will be returned to user.
   """
   before_tool_callback: Optional[BeforeToolCallback] = None
-  """Called before the tool is called.
+  """Callback or list of callbacks to be called before calling the tool.
+
+  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:
     tool: The tool to be called.
@@ -207,7 +255,10 @@ class LlmAgent(BaseAgent):
     the framework will skip calling the actual tool.
   """
   after_tool_callback: Optional[AfterToolCallback] = None
-  """Called after the tool is called.
+  """Callback or list of callbacks to be called after calling the tool.
+
+  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:
     tool: The tool to be called.
@@ -256,33 +307,119 @@ class LlmAgent(BaseAgent):
         ancestor_agent = ancestor_agent.parent_agent
       raise ValueError(f'No model found for {self.name}.')
 
-  def canonical_instruction(self, ctx: ReadonlyContext) -> str:
+  async def canonical_instruction(
+      self, ctx: ReadonlyContext
+  ) -> tuple[str, bool]:
     """The resolved self.instruction field to construct instruction for this agent.
 
     This method is only for use by Agent Development Kit.
+
+    Args:
+      ctx: The context to retrieve the session state.
+
+    Returns:
+      A tuple of (instruction, bypass_state_injection).
+      instruction: The resolved self.instruction field.
+      bypass_state_injection: Whether the instruction is based on
+      InstructionProvider.
     """
     if isinstance(self.instruction, str):
-      return self.instruction
+      return self.instruction, False
     else:
-      return self.instruction(ctx)
-
-  def canonical_global_instruction(self, ctx: ReadonlyContext) -> str:
+      instruction = self.instruction(ctx)
+      if inspect.isawaitable(instruction):
+        instruction = await instruction
+      return instruction, True
+
+  async def canonical_global_instruction(
+      self, ctx: ReadonlyContext
+  ) -> tuple[str, bool]:
     """The resolved self.instruction field to construct global instruction.
 
     This method is only for use by Agent Development Kit.
+
+    Args:
+      ctx: The context to retrieve the session state.
+
+    Returns:
+      A tuple of (instruction, bypass_state_injection).
+      instruction: The resolved self.global_instruction field.
+      bypass_state_injection: Whether the instruction is based on
+      InstructionProvider.
     """
     if isinstance(self.global_instruction, str):
-      return self.global_instruction
+      return self.global_instruction, False
     else:
-      return self.global_instruction(ctx)
+      global_instruction = self.global_instruction(ctx)
+      if inspect.isawaitable(global_instruction):
+        global_instruction = await global_instruction
+      return global_instruction, True
+
+  async def canonical_tools(
+      self, ctx: ReadonlyContext = None
+  ) -> list[BaseTool]:
+    """The resolved self.tools field as a list of BaseTool based on the context.
+
+    This method is only for use by Agent Development Kit.
+    """
+    resolved_tools = []
+    for tool_union in self.tools:
+      resolved_tools.extend(await _convert_tool_union_to_tools(tool_union, ctx))
+    return resolved_tools
+
+  @property
+  def canonical_before_model_callbacks(
+      self,
+  ) -> list[_SingleBeforeModelCallback]:
+    """The resolved self.before_model_callback field as a list of _SingleBeforeModelCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.before_model_callback:
+      return []
+    if isinstance(self.before_model_callback, list):
+      return self.before_model_callback
+    return [self.before_model_callback]
+
+  @property
+  def canonical_after_model_callbacks(self) -> list[_SingleAfterModelCallback]:
+    """The resolved self.after_model_callback field as a list of _SingleAfterModelCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.after_model_callback:
+      return []
+    if isinstance(self.after_model_callback, list):
+      return self.after_model_callback
+    return [self.after_model_callback]
+
+  @property
+  def canonical_before_tool_callbacks(
+      self,
+  ) -> list[BeforeToolCallback]:
+    """The resolved self.before_tool_callback field as a list of BeforeToolCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.before_tool_callback:
+      return []
+    if isinstance(self.before_tool_callback, list):
+      return self.before_tool_callback
+    return [self.before_tool_callback]
 
   @property
-  def canonical_tools(self) -> list[BaseTool]:
-    """The resolved self.tools field as a list of BaseTool.
+  def canonical_after_tool_callbacks(
+      self,
+  ) -> list[AfterToolCallback]:
+    """The resolved self.after_tool_callback field as a list of AfterToolCallback.
 
     This method is only for use by Agent Development Kit.
     """
-    return [_convert_tool_union_to_tool(tool) for tool in self.tools]
+    if not self.after_tool_callback:
+      return []
+    if isinstance(self.after_tool_callback, list):
+      return self.after_tool_callback
+    return [self.after_tool_callback]
 
   @property
   def _llm_flow(self) -> BaseLlmFlow:

google/adk/agents/loop_agent.py

@@ -58,5 +58,5 @@ class LoopAgent(BaseAgent):
   async def _run_live_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
-    raise NotImplementedError('The behavior for run_live is not defined yet.')
+    raise NotImplementedError('This is not supported yet for LoopAgent.')
     yield  # AsyncGenerator requires having at least one yield statement

google/adk/agents/parallel_agent.py

@@ -94,3 +94,10 @@ class ParallelAgent(BaseAgent):
     agent_runs = [agent.run_async(ctx) for agent in self.sub_agents]
     async for event in _merge_agent_run(agent_runs):
       yield event
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    raise NotImplementedError("This is not supported yet for ParallelAgent.")
+    yield  # AsyncGenerator requires having at least one yield statement

google/adk/agents/readonly_context.py

@@ -15,10 +15,11 @@
 from __future__ import annotations
 
 from types import MappingProxyType
-from typing import Any
+from typing import Any, Optional
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+  from google.genai import types
   from .invocation_context import InvocationContext
 
 
@@ -30,6 +31,11 @@ class ReadonlyContext:
   ) -> None:
     self._invocation_context = invocation_context
 
+  @property
+  def user_content(self) -> Optional[types.Content]:
+    """The user content that started this invocation. READONLY field."""
+    return self._invocation_context.user_content
+
   @property
   def invocation_id(self) -> str:
     """The current invocation id."""

google/adk/agents/run_config.py

@@ -22,7 +22,7 @@ from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import field_validator
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class StreamingMode(Enum):
@@ -37,6 +37,7 @@ class RunConfig(BaseModel):
   model_config = ConfigDict(
       extra='forbid',
   )
+  """The pydantic model config."""
 
   speech_config: Optional[types.SpeechConfig] = None
   """Speech configuration for the live agent."""
@@ -64,6 +65,9 @@ class RunConfig(BaseModel):
   output_audio_transcription: Optional[types.AudioTranscriptionConfig] = None
   """Output transcription for live agents with audio response."""
 
+  input_audio_transcription: Optional[types.AudioTranscriptionConfig] = None
+  """Input transcription for live agents with audio input from user."""
+
   max_llm_calls: int = 500
   """
   A limit on the total number of llm calls for a given run.

google/adk/agents/sequential_agent.py

@@ -23,6 +23,7 @@ from typing_extensions import override
 from ..agents.invocation_context import InvocationContext
 from ..events.event import Event
 from .base_agent import BaseAgent
+from .llm_agent import LlmAgent
 
 
 class SequentialAgent(BaseAgent):
@@ -40,6 +41,36 @@ class SequentialAgent(BaseAgent):
   async def _run_live_impl(
       self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
+    """Implementation for live SequentialAgent.
+
+    Compared to non-live case, live agents process a continous streams of audio
+    or video, so it doesn't have a way to tell if it's finished and should pass
+    to next agent or not. So we introduce a task_compelted() function so the
+    model can call this function to signal that it's finished the task and we
+    can move on to next agent.
+
+    Args:
+      ctx: The invocation context of the agent.
+    """
+    # There is no way to know if it's using live during init phase so we have to init it here
+    for sub_agent in self.sub_agents:
+      # add tool
+      def task_completed():
+        """
+        Signals that the model has successfully completed the user's question
+        or task.
+        """
+        return "Task completion signaled."
+
+      if isinstance(sub_agent, LlmAgent):
+        # Use function name to dedupe.
+        if task_completed.__name__ not in sub_agent.tools:
+          sub_agent.tools.append(task_completed)
+          sub_agent.instruction += f"""If you finished the user' request
+          according to its description, call {task_completed.__name__} function
+          to exit so the next agents can take over. When calling this function,
+          do not generate any text other than the function call.'"""
+
     for sub_agent in self.sub_agents:
       async for event in sub_agent.run_live(ctx):
         yield event

google/adk/agents/transcription_entry.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from typing import Union
 
 from google.genai import types
@@ -26,9 +27,11 @@ class TranscriptionEntry(BaseModel):
       arbitrary_types_allowed=True,
       extra='forbid',
   )
+  """The pydantic model config."""
 
-  role: str
-  """The role that created this data, typically "user" or "model"""
+  role: Optional[str] = None
+  """The role that created this data, typically "user" or "model". For function 
+  call, this is None."""
 
   data: Union[types.Blob, types.Content]
   """The data that can be used for transcription"""

google/adk/artifacts/base_artifact_service.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Abstract base class for artifact services."""
 
 from abc import ABC
 from abc import abstractmethod
@@ -25,7 +24,7 @@ class BaseArtifactService(ABC):
   """Abstract base class for artifact services."""
 
   @abstractmethod
-  def save_artifact(
+  async def save_artifact(
       self,
       *,
       app_name: str,
@@ -53,7 +52,7 @@ class BaseArtifactService(ABC):
     """
 
   @abstractmethod
-  def load_artifact(
+  async def load_artifact(
       self,
       *,
       app_name: str,
@@ -78,10 +77,9 @@ class BaseArtifactService(ABC):
     Returns:
       The artifact or None if not found.
     """
-    pass
 
   @abstractmethod
-  def list_artifact_keys(
+  async def list_artifact_keys(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> list[str]:
     """Lists all the artifact filenames within a session.
@@ -94,10 +92,9 @@ class BaseArtifactService(ABC):
     Returns:
         A list of all artifact filenames within a session.
     """
-    pass
 
   @abstractmethod
-  def delete_artifact(
+  async def delete_artifact(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> None:
     """Deletes an artifact.
@@ -108,10 +105,9 @@ class BaseArtifactService(ABC):
         session_id: The ID of the session.
         filename: The name of the artifact file.
     """
-    pass
 
   @abstractmethod
-  def list_versions(
+  async def list_versions(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> list[int]:
     """Lists all versions of an artifact.
@@ -125,4 +121,3 @@ class BaseArtifactService(ABC):
     Returns:
         A list of all available versions of the artifact.
     """
-    pass

google/adk/artifacts/gcs_artifact_service.py

@@ -23,7 +23,7 @@ from typing_extensions import override
 
 from .base_artifact_service import BaseArtifactService
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 
 class GcsArtifactService(BaseArtifactService):
@@ -77,7 +77,7 @@ class GcsArtifactService(BaseArtifactService):
     return f"{app_name}/{user_id}/{session_id}/{filename}/{version}"
 
   @override
-  def save_artifact(
+  async def save_artifact(
       self,
       *,
       app_name: str,
@@ -86,7 +86,7 @@ class GcsArtifactService(BaseArtifactService):
       filename: str,
       artifact: types.Part,
   ) -> int:
-    versions = self.list_versions(
+    versions = await self.list_versions(
         app_name=app_name,
         user_id=user_id,
         session_id=session_id,
@@ -107,7 +107,7 @@ class GcsArtifactService(BaseArtifactService):
     return version
 
   @override
-  def load_artifact(
+  async def load_artifact(
       self,
       *,
       app_name: str,
@@ -117,7 +117,7 @@ class GcsArtifactService(BaseArtifactService):
       version: Optional[int] = None,
   ) -> Optional[types.Part]:
     if version is None:
-      versions = self.list_versions(
+      versions = await self.list_versions(
           app_name=app_name,
           user_id=user_id,
           session_id=session_id,
@@ -141,7 +141,7 @@ class GcsArtifactService(BaseArtifactService):
     return artifact
 
   @override
-  def list_artifact_keys(
+  async def list_artifact_keys(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> list[str]:
     filenames = set()
@@ -165,10 +165,10 @@ class GcsArtifactService(BaseArtifactService):
     return sorted(list(filenames))
 
   @override
-  def delete_artifact(
+  async def delete_artifact(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> None:
-    versions = self.list_versions(
+    versions = await self.list_versions(
         app_name=app_name,
         user_id=user_id,
         session_id=session_id,
@@ -183,7 +183,7 @@ class GcsArtifactService(BaseArtifactService):
     return
 
   @override
-  def list_versions(
+  async def list_versions(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> list[int]:
     prefix = self._get_blob_name(app_name, user_id, session_id, filename, "")

google/adk/artifacts/in_memory_artifact_service.py

@@ -24,7 +24,7 @@ from typing_extensions import override
 
 from .base_artifact_service import BaseArtifactService
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 
 class InMemoryArtifactService(BaseArtifactService, BaseModel):
@@ -63,7 +63,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     return f"{app_name}/{user_id}/{session_id}/{filename}"
 
   @override
-  def save_artifact(
+  async def save_artifact(
       self,
       *,
       app_name: str,
@@ -80,7 +80,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     return version
 
   @override
-  def load_artifact(
+  async def load_artifact(
       self,
       *,
       app_name: str,
@@ -98,7 +98,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     return versions[version]
 
   @override
-  def list_artifact_keys(
+  async def list_artifact_keys(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> list[str]:
     session_prefix = f"{app_name}/{user_id}/{session_id}/"
@@ -114,7 +114,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     return sorted(filenames)
 
   @override
-  def delete_artifact(
+  async def delete_artifact(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> None:
     path = self._artifact_path(app_name, user_id, session_id, filename)
@@ -123,7 +123,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
     self.artifacts.pop(path, None)
 
   @override
-  def list_versions(
+  async def list_versions(
       self, *, app_name: str, user_id: str, session_id: str, filename: str
   ) -> list[int]:
     path = self._artifact_path(app_name, user_id, session_id, filename)

google/adk/auth/auth_credential.py

@@ -13,17 +13,21 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
+from pydantic import alias_generators
 from pydantic import BaseModel
+from pydantic import ConfigDict
 from pydantic import Field
 
 
 class BaseModelWithConfig(BaseModel):
-  model_config = {"extra": "allow"}
+  model_config = ConfigDict(
+      extra="allow",
+      alias_generator=alias_generators.to_camel,
+      populate_by_name=True,
+  )
+  """The pydantic model config."""
 
 
 class HttpCredentials(BaseModelWithConfig):