google-adk 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

google/adk/agents/llm_agent.py

@@ -0,0 +1,376 @@
+# 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 logging
+from typing import Any
+from typing import AsyncGenerator
+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
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+from typing_extensions import override
+from typing_extensions import TypeAlias
+
+from ..code_executors.base_code_executor import BaseCodeExecutor
+from ..events.event import Event
+from ..examples.base_example_provider import BaseExampleProvider
+from ..examples.example import Example
+from ..flows.llm_flows.auto_flow import AutoFlow
+from ..flows.llm_flows.base_llm_flow import BaseLlmFlow
+from ..flows.llm_flows.single_flow import SingleFlow
+from ..models.base_llm import BaseLlm
+from ..models.llm_request import LlmRequest
+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.function_tool import FunctionTool
+from ..tools.tool_context import ToolContext
+from .base_agent import BaseAgent
+from .callback_context import CallbackContext
+from .invocation_context import InvocationContext
+from .readonly_context import ReadonlyContext
+
+logger = logging.getLogger(__name__)
+
+
+BeforeModelCallback: TypeAlias = Callable[
+    [CallbackContext, LlmRequest], Optional[LlmResponse]
+]
+AfterModelCallback: TypeAlias = Callable[
+    [CallbackContext, LlmResponse],
+    Optional[LlmResponse],
+]
+BeforeToolCallback: TypeAlias = Callable[
+    [BaseTool, dict[str, Any], ToolContext],
+    Optional[dict],
+]
+AfterToolCallback: TypeAlias = Callable[
+    [BaseTool, dict[str, Any], ToolContext, dict],
+    Optional[dict],
+]
+
+InstructionProvider: TypeAlias = Callable[[ReadonlyContext], str]
+
+ToolUnion: TypeAlias = Union[Callable, BaseTool]
+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)
+  )
+
+
+class LlmAgent(BaseAgent):
+  """LLM-based Agent."""
+
+  model: Union[str, BaseLlm] = ''
+  """The model to use for the agent.
+
+  When not set, the agent will inherit the model from its ancestor.
+  """
+
+  instruction: Union[str, InstructionProvider] = ''
+  """Instructions for the LLM model, guiding the agent's behavior."""
+
+  global_instruction: Union[str, InstructionProvider] = ''
+  """Instructions for all the agents in the entire agent tree.
+
+  global_instruction ONLY takes effect in root agent.
+
+  For example: use global_instruction to make all agents have a stable identity
+  or personality.
+  """
+
+  tools: list[ToolUnion] = Field(default_factory=list)
+  """Tools available to this agent."""
+
+  generate_content_config: Optional[types.GenerateContentConfig] = None
+  """The additional content generation configurations.
+
+  NOTE: not all fields are usable, e.g. tools must be configured via `tools`,
+  thinking_config must be configured via `planner` in LlmAgent.
+
+  For example: use this config to adjust model temperature, configure safety
+  settings, etc.
+  """
+
+  # LLM-based agent transfer configs - Start
+  disallow_transfer_to_parent: bool = False
+  """Disallows LLM-controlled transferring to the parent agent."""
+  disallow_transfer_to_peers: bool = False
+  """Disallows LLM-controlled transferring to the peer agents."""
+  # LLM-based agent transfer configs - End
+
+  include_contents: Literal['default', 'none'] = 'default'
+  """Whether to include contents in the model request.
+
+  When set to 'none', the model request will not include any contents, such as
+  user messages, tool results, etc.
+  """
+
+  # Controlled input/output configurations - Start
+  input_schema: Optional[type[BaseModel]] = None
+  """The input schema when agent is used as a tool."""
+  output_schema: Optional[type[BaseModel]] = None
+  """The output schema when agent replies.
+
+  NOTE: when this is set, agent can ONLY reply and CANNOT use any tools, such as
+  function tools, RAGs, agent transfer, etc.
+  """
+  output_key: Optional[str] = None
+  """The key in session state to store the output of the agent.
+
+  Typically use cases:
+  - Extracts agent reply for later use, such as in tools, callbacks, etc.
+  - Connects agents to coordinate with each other.
+  """
+  # Controlled input/output configurations - End
+
+  # Advance features - Start
+  planner: Optional[BasePlanner] = None
+  """Instructs the agent to make a plan and execute it step by step.
+
+  NOTE: to use model's built-in thinking features, set the `thinking_config`
+  field in `google.adk.planners.built_in_planner`.
+
+  """
+
+  code_executor: Optional[BaseCodeExecutor] = None
+  """Allow agent to execute code blocks from model responses using the provided
+  CodeExecutor.
+
+  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.
+  """
+  # Advance features - End
+
+  # TODO: remove below fields after migration. - Start
+  # These fields are added back for easier migration.
+  examples: Optional[ExamplesUnion] = None
+  # TODO: remove above fields after migration. - End
+
+  # Callbacks - Start
+  before_model_callback: Optional[BeforeModelCallback] = None
+  """Called before calling the LLM.
+  Args:
+    callback_context: CallbackContext,
+    llm_request: LlmRequest, The raw model request. Callback can mutate the
+    request.
+
+  Returns:
+    The content to return to the user. When present, the model call will be
+    skipped and the provided content will be returned to user.
+  """
+  after_model_callback: Optional[AfterModelCallback] = None
+  """Called after calling LLM.
+
+  Args:
+    callback_context: CallbackContext,
+    llm_response: LlmResponse, the actual model response.
+
+  Returns:
+    The content to return to the user. When present, the actual model response
+    will be ignored and the provided content will be returned to user.
+  """
+  before_tool_callback: Optional[BeforeToolCallback] = None
+  """Called before the tool is called.
+
+  Args:
+    tool: The tool to be called.
+    args: The arguments to the tool.
+    tool_context: ToolContext,
+
+  Returns:
+    The tool response. When present, the returned tool response will be used and
+    the framework will skip calling the actual tool.
+  """
+  after_tool_callback: Optional[AfterToolCallback] = None
+  """Called after the tool is called.
+
+  Args:
+    tool: The tool to be called.
+    args: The arguments to the tool.
+    tool_context: ToolContext,
+    tool_response: The response from the tool.
+
+  Returns:
+    When present, the returned dict will be used as tool result.
+  """
+  # Callbacks - End
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    async for event in self._llm_flow.run_async(ctx):
+      self.__maybe_save_output_to_state(event)
+      yield event
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    async for event in self._llm_flow.run_live(ctx):
+      self.__maybe_save_output_to_state(event)
+      yield event
+    if ctx.end_invocation:
+      return
+
+  @property
+  def canonical_model(self) -> BaseLlm:
+    """The resolved self.model field as BaseLlm.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if isinstance(self.model, BaseLlm):
+      return self.model
+    elif self.model:  # model is non-empty str
+      return LLMRegistry.new_llm(self.model)
+    else:  # find model from ancestors.
+      ancestor_agent = self.parent_agent
+      while ancestor_agent is not None:
+        if isinstance(ancestor_agent, LlmAgent):
+          return ancestor_agent.canonical_model
+        ancestor_agent = ancestor_agent.parent_agent
+      raise ValueError(f'No model found for {self.name}.')
+
+  def canonical_instruction(self, ctx: ReadonlyContext) -> str:
+    """The resolved self.instruction field to construct instruction for this agent.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if isinstance(self.instruction, str):
+      return self.instruction
+    else:
+      return self.instruction(ctx)
+
+  def canonical_global_instruction(self, ctx: ReadonlyContext) -> str:
+    """The resolved self.instruction field to construct global instruction.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if isinstance(self.global_instruction, str):
+      return self.global_instruction
+    else:
+      return self.global_instruction(ctx)
+
+  @property
+  def canonical_tools(self) -> list[BaseTool]:
+    """The resolved self.tools field as a list of BaseTool.
+
+    This method is only for use by Agent Development Kit.
+    """
+    return [_convert_tool_union_to_tool(tool) for tool in self.tools]
+
+  @property
+  def _llm_flow(self) -> BaseLlmFlow:
+    if (
+        self.disallow_transfer_to_parent
+        and self.disallow_transfer_to_peers
+        and not self.sub_agents
+    ):
+      return SingleFlow()
+    else:
+      return AutoFlow()
+
+  def __maybe_save_output_to_state(self, event: Event):
+    """Saves the model output to state if needed."""
+    if (
+        self.output_key
+        and event.is_final_response()
+        and event.content
+        and event.content.parts
+    ):
+      result = ''.join(
+          [part.text if part.text else '' for part in event.content.parts]
+      )
+      if self.output_schema:
+        result = self.output_schema.model_validate_json(result).model_dump(
+            exclude_none=True
+        )
+      event.actions.state_delta[self.output_key] = result
+
+  @model_validator(mode='after')
+  def __model_validator_after(self) -> LlmAgent:
+    self.__check_output_schema()
+    return self
+
+  def __check_output_schema(self):
+    if not self.output_schema:
+      return
+
+    if (
+        not self.disallow_transfer_to_parent
+        or not self.disallow_transfer_to_peers
+    ):
+      logger.warning(
+          'Invalid config for agent %s: output_schema cannot co-exist with'
+          ' agent transfer configurations. Setting'
+          ' disallow_transfer_to_parent=True, disallow_transfer_to_peers=True',
+          self.name,
+      )
+      self.disallow_transfer_to_parent = True
+      self.disallow_transfer_to_peers = True
+
+    if self.sub_agents:
+      raise ValueError(
+          f'Invalid config for agent {self.name}: if output_schema is set,'
+          ' sub_agents must be empty to disable agent transfer.'
+      )
+
+    if self.tools:
+      raise ValueError(
+          f'Invalid config for agent {self.name}: if output_schema is set,'
+          ' tools must be empty'
+      )
+
+  @field_validator('generate_content_config', mode='after')
+  @classmethod
+  def __validate_generate_content_config(
+      cls, generate_content_config: Optional[types.GenerateContentConfig]
+  ) -> types.GenerateContentConfig:
+    if not generate_content_config:
+      return types.GenerateContentConfig()
+    if generate_content_config.thinking_config:
+      raise ValueError('Thinking config should be set via LlmAgent.planner.')
+    if generate_content_config.tools:
+      raise ValueError('All tools must be set via LlmAgent.tools.')
+    if generate_content_config.system_instruction:
+      raise ValueError(
+          'System instruction must be set via LlmAgent.instruction.'
+      )
+    if generate_content_config.response_schema:
+      raise ValueError(
+          'Response schema must be set via LlmAgent.output_schema.'
+      )
+    return generate_content_config
+
+
+Agent: TypeAlias = LlmAgent

google/adk/agents/loop_agent.py

@@ -0,0 +1,62 @@
+# 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.
+
+"""Loop agent implementation."""
+
+from __future__ import annotations
+
+from typing import AsyncGenerator
+from typing import Optional
+
+from typing_extensions import override
+
+from ..agents.invocation_context import InvocationContext
+from ..events.event import Event
+from .base_agent import BaseAgent
+
+
+class LoopAgent(BaseAgent):
+  """A shell agent that run its sub-agents in a loop.
+
+  When sub-agent generates an event with escalate or max_iterations are
+  reached, the loop agent will stop.
+  """
+
+  max_iterations: Optional[int] = None
+  """The maximum number of iterations to run the loop agent.
+
+  If not set, the loop agent will run indefinitely until a sub-agent
+  escalates.
+  """
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    times_looped = 0
+    while not self.max_iterations or times_looped < self.max_iterations:
+      for sub_agent in self.sub_agents:
+        async for event in sub_agent.run_async(ctx):
+          yield event
+          if event.actions.escalate:
+            return
+      times_looped += 1
+    return
+
+  @override
+  async def _run_live_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    raise NotImplementedError('The behavior for run_live is not defined yet.')
+    yield  # AsyncGenerator requires having at least one yield statement

google/adk/agents/parallel_agent.py

@@ -0,0 +1,96 @@
+# 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.
+
+"""Parallel agent implementation."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncGenerator
+
+from typing_extensions import override
+
+from ..agents.invocation_context import InvocationContext
+from ..events.event import Event
+from .base_agent import BaseAgent
+
+
+def _set_branch_for_current_agent(
+    current_agent: BaseAgent, invocation_context: InvocationContext
+):
+  invocation_context.branch = (
+      f"{invocation_context.branch}.{current_agent.name}"
+      if invocation_context.branch
+      else current_agent.name
+  )
+
+
+async def _merge_agent_run(
+    agent_runs: list[AsyncGenerator[Event, None]],
+) -> AsyncGenerator[Event, None]:
+  """Merges the agent run event generator.
+
+  This implementation guarantees for each agent, it won't move on until the
+  generated event is processed by upstream runner.
+
+  Args:
+      agent_runs: A list of async generators that yield events from each agent.
+
+  Yields:
+      Event: The next event from the merged generator.
+  """
+  tasks = [
+      asyncio.create_task(events_for_one_agent.__anext__())
+      for events_for_one_agent in agent_runs
+  ]
+  pending_tasks = set(tasks)
+
+  while pending_tasks:
+    done, pending_tasks = await asyncio.wait(
+        pending_tasks, return_when=asyncio.FIRST_COMPLETED
+    )
+    for task in done:
+      try:
+        yield task.result()
+
+        # Find the generator that produced this event and move it on.
+        for i, original_task in enumerate(tasks):
+          if task == original_task:
+            new_task = asyncio.create_task(agent_runs[i].__anext__())
+            tasks[i] = new_task
+            pending_tasks.add(new_task)
+            break  # stop iterating once found
+
+      except StopAsyncIteration:
+        continue
+
+
+class ParallelAgent(BaseAgent):
+  """A shell agent that run its sub-agents in parallel in isolated manner.
+
+  This approach is beneficial for scenarios requiring multiple perspectives or
+  attempts on a single task, such as:
+
+  - Running different algorithms simultaneously.
+  - Generating multiple responses for review by a subsequent evaluation agent.
+  """
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    _set_branch_for_current_agent(self, ctx)
+    agent_runs = [agent.run_async(ctx) for agent in self.sub_agents]
+    async for event in _merge_agent_run(agent_runs):
+      yield event

google/adk/agents/readonly_context.py

@@ -0,0 +1,46 @@
+# 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 types import MappingProxyType
+from typing import Any
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from .invocation_context import InvocationContext
+
+
+class ReadonlyContext:
+
+  def __init__(
+      self,
+      invocation_context: InvocationContext,
+  ) -> None:
+    self._invocation_context = invocation_context
+
+  @property
+  def invocation_id(self) -> str:
+    """The current invocation id."""
+    return self._invocation_context.invocation_id
+
+  @property
+  def agent_name(self) -> str:
+    """The name of the agent that is currently running."""
+    return self._invocation_context.agent.name
+
+  @property
+  def state(self) -> MappingProxyType[str, Any]:
+    """The state of the current session. READONLY field."""
+    return MappingProxyType(self._invocation_context.session.state)

google/adk/agents/remote_agent.py

@@ -0,0 +1,50 @@
+# 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.
+
+import json
+from typing import AsyncGenerator
+
+from pydantic import Field
+import requests
+from typing_extensions import override
+
+from ..events.event import Event
+from .base_agent import BaseAgent
+from .invocation_context import InvocationContext
+
+
+class RemoteAgent(BaseAgent):
+  """Experimental, do not use."""
+
+  url: str
+
+  sub_agents: list[BaseAgent] = Field(
+      default_factory=list, init=False, frozen=True
+  )
+  """Sub-agent is dsiabled in RemoteAgent."""
+
+  @override
+  async def _run_async_impl(
+      self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    data = {
+        'invocation_id': ctx.invocation_id,
+        'session': ctx.session.model_dump(exclude_none=True),
+    }
+    events = requests.post(self.url, data=json.dumps(data), timeout=120)
+    events.raise_for_status()
+    for event in events.json():
+      e = Event.model_validate(event)
+      e.author = self.name
+      yield e