google--adk-python/contributing/samples/integrations/gepa/adk_agent.py

# Copyright 2026 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.

"""ADK utils for a LLMAgent interacting with a simulation environment."""

from __future__ import annotations

import asyncio
from collections.abc import Generator
from typing import Any
from typing import Dict
from typing import Optional
from typing import Protocol
from typing import runtime_checkable

from absl import logging
from google.adk import runners
from google.adk.agents import base_agent
from google.adk.agents import llm_agent
from google.adk.agents import loop_agent
from google.adk.events import event as event_lib
from google.adk.models import google_llm
from google.adk.planners import built_in_planner
from google.adk.tools import base_tool
from google.genai import types
from retry import api as retry


class EnvResponse(Protocol):
  """Environment response protocol."""

  observation: str
  done: bool
  reward: float


@runtime_checkable
class Env(Protocol):
  """Environment protocol."""

  def step(self, action: types.Part) -> EnvResponse:
    """Steps the environment with the given action."""
    ...

  def reset(self, task_index: int) -> EnvResponse:
    """Resets the environment to the given task index."""
    ...


class _Tool(base_tool.BaseTool):
  """A tool that executes an action in the environment."""

  class Config:
    arbitrary_types_allowed = True

  def __init__(
      self,
      function_declaration: types.FunctionDeclaration,
      env: Env,
  ):
    """Initializes the tool.

    Args:
      function_declaration: The function declaration of the tool.
      env: The environment to interact with.
    """
    super().__init__(
        name=function_declaration.name,
        description=function_declaration.description,
    )
    self._function_declaration = function_declaration
    self._env = env

  def _get_declaration(self) -> types.FunctionDeclaration:
    return self._function_declaration

  async def run_async(self, *, args: Dict[str, Any], tool_context: Any) -> str:
    """Runs the tool by converting tool call to env action and stepping env."""
    env_response = self._env.step(
        types.Part(function_call=types.FunctionCall(name=self.name, args=args))
    )
    # We modify the ADK session state with the updates from the environment,
    # in particular `done` and `reward`. These can be consumed downstream for
    # instance to extract the trajectory reward or interrupt the loop.
    tool_context.actions.state_delta['done'] = env_response.done
    tool_context.actions.state_delta['reward'] = env_response.reward
    tool_context.actions.skip_summarization = True
    if env_response.done:
      tool_context.actions.escalate = True
    return env_response.observation


def _default_retry_options() -> types.HttpRetryOptions:
  return types.HttpRetryOptions(
      initial_delay=2,
      attempts=4,
      max_delay=None,
      exp_base=2.0,
  )


def _adk_agent(
    instruction: str,
    tools: list[base_tool.BaseTool],
    temperature: float,
    model: str | None = None,
    name: str | None = None,
) -> llm_agent.LlmAgent:
  """Creates an ADK LLM agent with the given instruction and tools.

  Args:
    instruction: The instruction for the agent.
    tools: The tools for the agent to use.
    temperature: The temperature for the LLM.
    model: Model to use with the ADK LLMAgent ; defaults to `gemini-2.5-flash`.
    name: Name to set for the ADK LLM agent.

  Returns:
    An ADK LLM agent.
  """
  # TDOO - Allow more flexibility in configuring the agent used in the loop.
  return llm_agent.LlmAgent(
      name=name or 'agent',
      model=google_llm.Gemini(
          model=model or 'gemini-2.5-flash',
          retry_options=_default_retry_options(),
      ),
      planner=built_in_planner.BuiltInPlanner(
          thinking_config=types.ThinkingConfig(
              thinking_budget=-1, include_thoughts=False
          )
      ),
      instruction=instruction,
      tools=tools,
      generate_content_config=types.GenerateContentConfig(
          temperature=temperature,
          tool_config=types.ToolConfig(
              function_calling_config=types.FunctionCallingConfig(
                  mode=types.FunctionCallingConfigMode.VALIDATED
              )
          ),
          http_options=types.HttpOptions(
              timeout=30000,
              retry_options=_default_retry_options(),
          ),
      ),
  )


class _UserAgent(base_agent.BaseAgent):
  """An agent that wraps the provided environment and simulates a user."""

  env: Env

  class Config:
    arbitrary_types_allowed = True

  async def _run_async_impl(self, ctx: Any) -> Any:
    """Runs the user agent."""
    if not ctx.session.events:
      raise ValueError(
          'No prior session events, this is unexpected as the user agent cannot'
          ' be the first step in the interaction loop.'
      )
    last_event = ctx.session.events[-1]

    # Function tool
    if last_event.content and last_event.content.role == 'user':
      return

    if last_event.content and last_event.content.parts:
      next_message = '\n\n'.join([p.text for p in last_event.content.parts])
    else:
      logging.warn('Empty content with event=%s', last_event)
      next_message = ''
    env_response = retry.retry_call(
        self.env.step,
        fargs=(types.Part(text=next_message),),
        tries=3,
        delay=2,
        backoff=2,
    )

    output_event = event_lib.Event(
        content=types.Content(
            parts=[types.Part(text=env_response.observation)], role='user'
        ),
        author='user',
    )
    if env_response.done:
      output_event.actions.escalate = True
    output_event.actions.state_delta['reward'] = env_response.reward
    output_event.actions.state_delta['done'] = env_response.done
    yield output_event


def run_environment_loop(
    instruction: str,
    env: Env,
    temperature: float,
    tools: list[types.FunctionDeclaration],
    task_index: int,
    max_num_steps: int = 30,
    plugins: Optional[Any] = None,
    agent_model: str | None = None,
    agent_name: str | None = None,
) -> Generator[event_lib.Event]:
  """Defines and runs an ADK LLM Agent in the provided simulation environment.

  Args:
    instruction: The instruction for the agent.
    env: The environment to interact with.
    temperature: The temperature for the LLM.
    tools: The tools for the agent to use.
    task_index: The index of the task to run.
    max_num_steps: The maximum number of steps to run LLM agent - environment
      interaction loop.
    plugins: Optional plugins to use in the runner.
    agent_model: Model to use with the ADK LLMAgent ; defaults to
      `gemini-2.5-flash`.
    agent_name: Name to set for the ADK LLM agent.

  Returns:
    A generator of events from the agent run.

  Yields:
    All the events from the environment loop including:
      - Initial message from environment reset
      - LLMAgent generated text and function calls
      - Environment tools / users generated text responses
      - Environment user
  """
  # We use an agent loop to orchestrate the llm-agent and the environment
  # interactions. In particular to:
  # - ensure that LLMAgent and environment / user are called one after the
  # other
  # - the number of interaction steps is pre-defined (early exit is possible).
  agent = loop_agent.LoopAgent(
      name='env_loop_agent',
      max_iterations=max_num_steps,
      sub_agents=[
          _adk_agent(
              instruction=instruction,
              tools=[_Tool(t, env) for t in tools],
              temperature=temperature,
              model=agent_model,
              name=agent_name,
          ),
          _UserAgent(
              name='user_agent',
              env=env,
          ),
      ],
  )

  async def _async_run():
    runner = runners.InMemoryRunner(
        agent=agent,
        app_name='eval_app',
        plugins=plugins,
    )
    session = await runner.session_service.create_session(
        app_name='eval_app', user_id='eval_user'
    )
    env_reset_res = env.reset(task_index=task_index)
    initial_message = types.Content(
        role='user', parts=[types.Part(text=env_reset_res.observation)]
    )
    # The initial message is generated by the environment `reset` within the
    # implementation of this function - as the first step of the trace.
    # We yield this first step to ensure we provide a full trace to the user.
    events = [
        event_lib.Event(
            author='user',
            content=initial_message,
        )
    ]
    async for event in runner.run_async(
        user_id=session.user_id,
        session_id=session.id,
        new_message=initial_message,
    ):
      events.append(event)
    return events

  return asyncio.run(_async_run())