mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Add util to run ADK LLM Agent with simulation environment
PiperOrigin-RevId: 825674499
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
e7f7705eba
commit
87f415a7c3
@@ -0,0 +1,4 @@
|
||||
# Example: optimizing an ADK agent with Genetic-Pareto
|
||||
|
||||
This directory contains an example demonstrating how to use the Agent Development
|
||||
Kit (ADK) to run and optimize an LLM-based agent in a simulated environment with the Genetic-Pareto prompt optimization algorithm ([GEPA: Reflective Prompt Evolution Can Outperform Reinforcement Learning](https://arxiv.org/abs/2507.19457)) on benchmarks like Tau-bench.
|
||||
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,262 @@
|
||||
# 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.
|
||||
|
||||
"""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.tools import base_tool
|
||||
from google.genai import types
|
||||
|
||||
|
||||
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 _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=model or 'gemini-2.5-flash',
|
||||
instruction=instruction,
|
||||
tools=tools,
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
tool_config=types.ToolConfig(
|
||||
function_calling_config=types.FunctionCallingConfig(
|
||||
mode=types.FunctionCallingConfigMode.VALIDATED
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _UserAgent(base_agent.BaseAgent):
|
||||
"""An agent that wraps the provided environment and simulates an 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 = last_event.content.parts[-1].text
|
||||
else:
|
||||
logging.warn('Empty content with event=%s', last_event)
|
||||
next_message = ''
|
||||
env_response = self.env.step(types.Part(text=next_message))
|
||||
|
||||
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,
|
||||
),
|
||||
],
|
||||
)
|
||||
runner = runners.InMemoryRunner(
|
||||
agent=agent,
|
||||
app_name='eval_app',
|
||||
plugins=plugins,
|
||||
)
|
||||
session = asyncio.run(
|
||||
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.
|
||||
yield event_lib.Event(
|
||||
author='user',
|
||||
content=initial_message,
|
||||
)
|
||||
for event in runner.run(
|
||||
user_id=session.user_id,
|
||||
session_id=session.id,
|
||||
new_message=initial_message,
|
||||
):
|
||||
yield event
|
||||
@@ -0,0 +1,348 @@
|
||||
# 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 asyncio
|
||||
import dataclasses
|
||||
from unittest import mock
|
||||
|
||||
from gepa import adk_agent
|
||||
from google.adk import runners
|
||||
from google.adk.agents import base_agent
|
||||
from google.adk.events import event as event_lib
|
||||
from google.adk.plugins import base_plugin
|
||||
from google.genai import types
|
||||
|
||||
|
||||
class _TestPlugin(base_plugin.BasePlugin):
|
||||
|
||||
def __init__(self, outputs):
|
||||
super().__init__(name="test-pluggin")
|
||||
self._model_output_idx = 0
|
||||
self.got_llm_requests = []
|
||||
self._outputs = outputs
|
||||
|
||||
async def before_model_callback(self, *, callback_context, llm_request):
|
||||
self.got_llm_requests.append(llm_request)
|
||||
if self._model_output_idx < len(self._outputs):
|
||||
out = self._outputs[self._model_output_idx]
|
||||
self._model_output_idx += 1
|
||||
return out
|
||||
return event_lib.Event(
|
||||
error_code="empty test list",
|
||||
author="agent",
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class EnvResponse:
|
||||
observation: str
|
||||
done: bool
|
||||
reward: float
|
||||
|
||||
|
||||
class _TestEnv:
|
||||
|
||||
def __init__(self, responses):
|
||||
self._responses = responses
|
||||
self._idx = 0
|
||||
|
||||
def step(self, action):
|
||||
del action
|
||||
if self._idx < len(self._responses):
|
||||
resp = self._responses[self._idx]
|
||||
self._idx += 1
|
||||
else:
|
||||
resp = EnvResponse("out-of-bound", done=True, reward=0)
|
||||
return resp
|
||||
|
||||
def reset(self, task_index: int):
|
||||
del task_index
|
||||
return EnvResponse("reset-obs", done=False, reward=42)
|
||||
|
||||
|
||||
def test_default_flow():
|
||||
model_outputs = [
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="ab")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name="test_tool",
|
||||
args=dict(tool_inputs="fake-tool-inputs"),
|
||||
)
|
||||
)
|
||||
],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="cd")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
]
|
||||
events = adk_agent.run_environment_loop(
|
||||
instruction="some-instruction",
|
||||
env=_TestEnv([
|
||||
EnvResponse("some-obs-1", done=False, reward=123),
|
||||
EnvResponse("tool-response", done=False, reward=45),
|
||||
EnvResponse("some-obs-2", done=False, reward=67),
|
||||
]),
|
||||
temperature=0,
|
||||
tools=[
|
||||
types.FunctionDeclaration(
|
||||
name="test_tool",
|
||||
description="test_tool",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tool_inputs": {
|
||||
"type": "string",
|
||||
"description": "tool_inputs",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
],
|
||||
task_index=0,
|
||||
max_num_steps=3,
|
||||
plugins=[
|
||||
_TestPlugin(model_outputs),
|
||||
],
|
||||
)
|
||||
events = list(events)
|
||||
want = [
|
||||
"reset-obs",
|
||||
"ab",
|
||||
"some-obs-1",
|
||||
"test_tool",
|
||||
"tool-response",
|
||||
"cd",
|
||||
"some-obs-2",
|
||||
]
|
||||
|
||||
def _extract_from_event(event):
|
||||
if not event.content:
|
||||
return ""
|
||||
if len(event.content.parts) != 1:
|
||||
return ""
|
||||
part = event.content.parts[0]
|
||||
if part.function_call:
|
||||
return part.function_call.name
|
||||
if part.function_response:
|
||||
return part.function_response.response.get("result")
|
||||
return part.text
|
||||
|
||||
got = [_extract_from_event(e) for e in events]
|
||||
assert got == want
|
||||
|
||||
got_rewards = [e.actions.state_delta.get("reward") for e in events]
|
||||
assert got_rewards == [None, None, 123, None, 45, None, 67]
|
||||
|
||||
|
||||
def test_intermediary_step_is_done():
|
||||
model_outputs = [
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="ab")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="cd")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
]
|
||||
events = adk_agent.run_environment_loop(
|
||||
instruction="some-instruction",
|
||||
env=_TestEnv([
|
||||
EnvResponse("some-obs-1", done=True, reward=0),
|
||||
EnvResponse("some-obs-2", done=False, reward=0),
|
||||
]),
|
||||
temperature=0,
|
||||
tools=[],
|
||||
task_index=0,
|
||||
max_num_steps=5,
|
||||
plugins=[
|
||||
_TestPlugin(model_outputs),
|
||||
],
|
||||
)
|
||||
want_text = ["reset-obs", "ab", "some-obs-1"]
|
||||
got = [e.content.parts[0].text for e in events]
|
||||
assert got == want_text
|
||||
|
||||
|
||||
def test_intermediary_tool_step_is_done():
|
||||
model_outputs = [
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="ab")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name="test_tool",
|
||||
args=dict(tool_inputs="fake-tool-inputs"),
|
||||
)
|
||||
)
|
||||
],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="cd")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
]
|
||||
events = adk_agent.run_environment_loop(
|
||||
instruction="some-instruction",
|
||||
env=_TestEnv([
|
||||
EnvResponse("some-obs-1", done=False, reward=123),
|
||||
EnvResponse("tool-response", done=True, reward=45),
|
||||
EnvResponse("some-obs-2", done=False, reward=67),
|
||||
]),
|
||||
temperature=0,
|
||||
tools=[
|
||||
types.FunctionDeclaration(
|
||||
name="test_tool",
|
||||
description="test_tool",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tool_inputs": {
|
||||
"type": "string",
|
||||
"description": "tool_inputs",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
],
|
||||
task_index=0,
|
||||
max_num_steps=3,
|
||||
plugins=[
|
||||
_TestPlugin(model_outputs),
|
||||
],
|
||||
)
|
||||
events = list(events)
|
||||
want = ["reset-obs", "ab", "some-obs-1", "test_tool", "tool-response"]
|
||||
|
||||
def _extract_from_event(event):
|
||||
if not event.content:
|
||||
return ""
|
||||
if len(event.content.parts) != 1:
|
||||
return ""
|
||||
part = event.content.parts[0]
|
||||
if part.function_call:
|
||||
return part.function_call.name
|
||||
if part.function_response:
|
||||
return part.function_response.response.get("result")
|
||||
return part.text
|
||||
|
||||
got = [_extract_from_event(e) for e in events]
|
||||
assert got == want
|
||||
|
||||
|
||||
def test_llm_request():
|
||||
model_outputs = [
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="ab")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
event_lib.Event(
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="cd")],
|
||||
role="model",
|
||||
),
|
||||
author="agent",
|
||||
),
|
||||
]
|
||||
test_plugin = _TestPlugin(model_outputs)
|
||||
events = adk_agent.run_environment_loop(
|
||||
instruction="some-instruction",
|
||||
env=_TestEnv([
|
||||
EnvResponse("some-obs-1", done=False, reward=123),
|
||||
EnvResponse("some-obs-2", done=False, reward=67),
|
||||
]),
|
||||
temperature=0.123,
|
||||
tools=[],
|
||||
task_index=0,
|
||||
max_num_steps=2,
|
||||
plugins=[test_plugin],
|
||||
)
|
||||
_ = list(events)
|
||||
|
||||
assert len(test_plugin.got_llm_requests) == 2
|
||||
got = test_plugin.got_llm_requests[-1]
|
||||
assert "some-instruction" in got.config.system_instruction
|
||||
assert got.config.temperature == 0.123
|
||||
got_parts = [c.parts[0].text for c in got.contents]
|
||||
assert got_parts == ["reset-obs", "ab", "some-obs-1"]
|
||||
|
||||
|
||||
def test_model_name_is_set():
|
||||
class _MockAgent(base_agent.BaseAgent):
|
||||
|
||||
async def _run_async_impl(self, ctx):
|
||||
pass
|
||||
|
||||
async def _mock_create_session(*args, **kwargs):
|
||||
del args, kwargs
|
||||
await asyncio.sleep(0.1)
|
||||
return
|
||||
|
||||
with mock.patch.object(runners, "InMemoryRunner") as mock_runner_cls:
|
||||
mock_runner = mock_runner_cls.return_value
|
||||
mock_runner.session_service.create_session.side_effect = (
|
||||
_mock_create_session
|
||||
)
|
||||
mock_runner.run.return_value = []
|
||||
next(
|
||||
adk_agent.run_environment_loop(
|
||||
instruction="some-instruction",
|
||||
env=_TestEnv([]),
|
||||
temperature=0.123,
|
||||
tools=[],
|
||||
task_index=0,
|
||||
agent_model="some-test-model",
|
||||
plugins=[_TestPlugin([])],
|
||||
)
|
||||
)
|
||||
mock_runner_cls.assert_called_once()
|
||||
_, runner_kwargs = mock_runner_cls.call_args
|
||||
assert runner_kwargs["agent"].sub_agents[0].model == "some-test-model"
|
||||
Reference in New Issue
Block a user