From 4e3e2cb58858e08a79bc6119ad49b6c049dbc0d0 Mon Sep 17 00:00:00 2001 From: Keyur Joshi Date: Tue, 3 Mar 2026 10:16:06 -0800 Subject: [PATCH] feat: Add GEPA root agent prompt optimizer details: * Uses GEPA (https://gepa-ai.github.io/gepa/) to optimize the instructions for the root agent. Can be extended to sub-agents and other components in the future. * GEPA package is imported dynamically; you do not need to install it along with ADK unless you plan to use this optimizer. Co-authored-by: Keyur Joshi PiperOrigin-RevId: 878009649 --- pyproject.toml | 1 + .../gepa_root_agent_prompt_optimizer.py | 323 ++++++++++++++++++ .../gepa_root_agent_prompt_optimizer_test.py | 264 ++++++++++++++ 3 files changed, 588 insertions(+) create mode 100644 src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py create mode 100644 tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py diff --git a/pyproject.toml b/pyproject.toml index 0441c72d..d0f3cd94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -109,6 +109,7 @@ community = [ eval = [ # go/keep-sorted start "Jinja2>=3.1.4,<4.0.0", # For eval template rendering + "gepa>=0.1.0", "google-cloud-aiplatform[evaluation]>=1.100.0", "pandas>=2.2.3", "rouge-score>=0.1.2", diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py new file mode 100644 index 00000000..0627aced --- /dev/null +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -0,0 +1,323 @@ +# 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. + +from __future__ import annotations + +import asyncio +import logging +from typing import Any +from typing import Optional + +from google.genai import types as genai_types +from pydantic import BaseModel +from pydantic import Field + +from ..agents.llm_agent import Agent +from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE +from ..models.llm_request import LlmRequest +from ..models.llm_response import LlmResponse +from ..models.registry import LLMRegistry +from ..utils.context_utils import Aclosing +from ..utils.feature_decorator import experimental +from .agent_optimizer import AgentOptimizer +from .data_types import BaseAgentWithScores +from .data_types import OptimizerResult +from .data_types import UnstructuredSamplingResult +from .sampler import Sampler + +_logger = logging.getLogger("google_adk." + __name__) + +_AGENT_PROMPT_NAME = "agent_prompt" + + +class GEPARootAgentPromptOptimizerConfig(BaseModel): + """Contains configuration options required by the GEPARootAgentPromptOptimizer.""" + + optimizer_model: str = Field( + default="gemini-2.5-flash", + description=( + "The model used to analyze the eval results and optimize the agent." + ), + ) + + model_configuration: genai_types.GenerateContentConfig = Field( + default_factory=lambda: genai_types.GenerateContentConfig( + thinking_config=genai_types.ThinkingConfig( + include_thoughts=True, + thinking_budget=10240, + ) + ), + description="The configuration for the optimizer model.", + ) + + max_metric_calls: int = Field( + default=100, + description="The maximum number of metric calls (evaluations) to make.", + ) + + reflection_minibatch_size: int = Field( + default=3, + description="The number of examples to use for reflection.", + ) + + run_dir: Optional[str] = Field( + default=None, + description=( + "The directory to save the intermediate/final optimization results." + ), + ) + + +class GEPARootAgentPromptOptimizerResult(OptimizerResult[BaseAgentWithScores]): + """The final result of the GEPARootAgentPromptOptimizer.""" + + gepa_result: Optional[dict[str, Any]] = Field( + default=None, + description="The raw result dictionary from the GEPA optimizer.", + ) + + +def _create_agent_gepa_adapter_class(): + """Creates the _AgentGEPAAdapter class dynamically to avoid top-level gepa imports.""" + from gepa.core.adapter import EvaluationBatch + from gepa.core.adapter import GEPAAdapter + + class _AgentGEPAAdapter(GEPAAdapter[str, dict[str, Any], dict[str, Any]]): + """A GEPA adapter for ADK agents.""" + + def __init__( + self, + initial_agent: Agent, + sampler: Sampler[UnstructuredSamplingResult], + main_loop: asyncio.AbstractEventLoop, + ): + self._initial_agent = initial_agent + self._sampler = sampler + self._main_loop = main_loop + + self._train_example_ids = set(sampler.get_train_example_ids()) + self._validation_example_ids = set(sampler.get_validation_example_ids()) + + def evaluate( + self, + batch: list[str], + candidate: dict[str, str], + capture_traces: bool = False, + ) -> EvaluationBatch[dict[str, Any], dict[str, Any]]: + prompt = candidate[_AGENT_PROMPT_NAME] + _logger.info( + "Evaluating agent on batch:\n%s\nwith prompt:\n%s", batch, prompt + ) + # Clone the agent and update the instruction + new_agent = self._initial_agent.clone(update={"instruction": prompt}) + + if set(batch) <= self._train_example_ids: + example_set = "train" + elif set(batch) <= self._validation_example_ids: + example_set = "validation" + else: + raise ValueError(f"Invalid batch composition: {batch}") + + # Run the evaluation in the main loop + future = asyncio.run_coroutine_threadsafe( + self._sampler.sample_and_score( + new_agent, + example_set=example_set, + batch=batch, + capture_full_eval_data=capture_traces, + ), + self._main_loop, + ) + result: UnstructuredSamplingResult = future.result() + + scores = [] + outputs = [] + trajectories = [] + + for example_id in batch: + score = result.scores[example_id] + scores.append(score) + + eval_data = result.data.get(example_id, {}) if result.data else {} + outputs.append(eval_data) + trajectories.append(eval_data) + + return EvaluationBatch( + outputs=outputs, scores=scores, trajectories=trajectories + ) + + def make_reflective_dataset( + self, + candidate: dict[str, str], + eval_batch: EvaluationBatch[dict[str, Any], dict[str, Any]], + components_to_update: list[str], + ) -> dict[str, list[dict[str, Any]]]: + dataset: list[dict[str, Any]] = [] + trace_instances: list[tuple[float, dict[str, Any]]] = list( + zip( + eval_batch.scores, + eval_batch.trajectories, + strict=True, + ) + ) + for trace_instance in trace_instances: + score, eval_data = trace_instance + + dataset.append({ + _AGENT_PROMPT_NAME: candidate[_AGENT_PROMPT_NAME], + "score": score, + "eval_data": eval_data, + }) + + # same data for all components (should be only one) + result = {comp: dataset for comp in components_to_update} + + return result + + return _AgentGEPAAdapter + + +@experimental +class GEPARootAgentPromptOptimizer( + AgentOptimizer[UnstructuredSamplingResult, BaseAgentWithScores] +): + """An optimizer that improves the root agent prompt using the GEPA framework.""" + + def __init__( + self, + config: GEPARootAgentPromptOptimizerConfig, + ): + self._config = config + llm_registry = LLMRegistry() + self._llm_class = llm_registry.resolve(self._config.optimizer_model) + + async def optimize( + self, + initial_agent: Agent, + sampler: Sampler[UnstructuredSamplingResult], + ) -> GEPARootAgentPromptOptimizerResult: + """Runs the GEPARootAgentPromptOptimizer. + + Args: + initial_agent: The initial agent whose prompt is to be optimized. Only the + root agent prompt will be optimized. + sampler: The interface used to get training and validation example UIDs, + request agent evaluations, and get useful data for optimizing the agent. + + Returns: + The final result of the optimization process, containing the optimized + agent instance, its scores on the validation examples, and other metrics. + """ + if initial_agent.sub_agents: + _logger.warning( + "The GEPARootAgentPromptOptimizer will not optimize prompts for" + " sub-agents." + ) + + _logger.info("Setting up the GEPA optimizer...") + + try: + import gepa # lazy import as gepa is not in core ADK package + + _AgentGEPAAdapter = _create_agent_gepa_adapter_class() + except ImportError as e: + raise ImportError(MISSING_EVAL_DEPENDENCIES_MESSAGE) from e + + loop = asyncio.get_running_loop() + + adapter = _AgentGEPAAdapter( + initial_agent=initial_agent, + sampler=sampler, + main_loop=loop, + ) + + llm = self._llm_class(model=self._config.optimizer_model) + + def reflection_lm(prompt: str) -> str: + llm_request = LlmRequest( + model=self._config.optimizer_model, + config=self._config.model_configuration, + contents=[ + genai_types.Content( + parts=[genai_types.Part(text=prompt)], + role="user", + ) + ], + ) + + async def _generate(): + response_text = "" + async with Aclosing(llm.generate_content_async(llm_request)) as agen: + async for llm_response in agen: + llm_response: LlmResponse + generated_content: genai_types.Content = llm_response.content + if not generated_content.parts: + continue + response_text = "".join( + part.text + for part in generated_content.parts + if part.text and not part.thought + ) + return response_text + + future = asyncio.run_coroutine_threadsafe(_generate(), loop) + return future.result() + + train_ids = sampler.get_train_example_ids() + val_ids = sampler.get_validation_example_ids() + + if set(train_ids).intersection(val_ids): + _logger.warning( + "The training and validation example UIDs overlap. This WILL cause" + " aliasing issues unless each common UID refers to the same example" + " in both sets." + ) + + def run_gepa(): + return gepa.optimize( + seed_candidate={_AGENT_PROMPT_NAME: initial_agent.instruction}, + trainset=train_ids, + valset=val_ids, + adapter=adapter, + max_metric_calls=self._config.max_metric_calls, + reflection_lm=reflection_lm, + reflection_minibatch_size=self._config.reflection_minibatch_size, + run_dir=self._config.run_dir, + ) + + _logger.info("Running the GEPA optimizer...") + + gepa_results = await loop.run_in_executor(None, run_gepa) + + _logger.info("GEPA optimization finished. Preparing final results...") + + optimized_prompts = [ + candidate[_AGENT_PROMPT_NAME] for candidate in gepa_results.candidates + ] + scores = gepa_results.val_aggregate_scores + + optimized_agents = [ + BaseAgentWithScores( + optimized_agent=initial_agent.clone( + update={"instruction": optimized_prompt}, + ), + overall_score=score, + ) + for optimized_prompt, score in zip(optimized_prompts, scores) + ] + + return GEPARootAgentPromptOptimizerResult( + optimized_agents=optimized_agents, + gepa_result=gepa_results.to_dict(), + ) diff --git a/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py b/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py new file mode 100644 index 00000000..bd5da524 --- /dev/null +++ b/tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py @@ -0,0 +1,264 @@ +# 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. + +from __future__ import annotations + +import asyncio +import sys + +from google.adk.agents.llm_agent import Agent +from google.adk.optimization.data_types import UnstructuredSamplingResult +from google.adk.optimization.gepa_root_agent_prompt_optimizer import _create_agent_gepa_adapter_class +from google.adk.optimization.gepa_root_agent_prompt_optimizer import GEPARootAgentPromptOptimizer +from google.adk.optimization.gepa_root_agent_prompt_optimizer import GEPARootAgentPromptOptimizerConfig +from google.adk.optimization.sampler import Sampler +import pytest + + +class MockEvaluationBatch: + + def __init__(self, outputs, scores, trajectories): + self.outputs = outputs + self.scores = scores + self.trajectories = trajectories + + +class MockGEPAAdapter: + """Mock that supports generic type hints.""" + + def __class_getitem__(cls, item): + return cls + + +@pytest.fixture(name="mock_gepa") +def fixture_mock_gepa(mocker): + # mock gepa before it gets imported by the optimizer module + mock_gepa_module = mocker.MagicMock() + mock_gepa_adapter = mocker.MagicMock() + + mock_gepa_adapter.EvaluationBatch = MockEvaluationBatch + mock_gepa_adapter.GEPAAdapter = MockGEPAAdapter + + mock_gepa_module.core = mocker.MagicMock() + mock_gepa_module.core.adapter = mock_gepa_adapter + + mocker.patch.dict( + sys.modules, + { + "gepa": mock_gepa_module, + "gepa.core": mock_gepa_module.core, + "gepa.core.adapter": mock_gepa_adapter, + }, + ) + return mock_gepa_module + + +@pytest.fixture +def mock_sampler(mocker): + sampler = mocker.MagicMock(spec=Sampler) + sampler.get_train_example_ids.return_value = ["train1", "train2"] + sampler.get_validation_example_ids.return_value = ["val1", "val2"] + return sampler + + +@pytest.fixture +def mock_agent(mocker): + agent = mocker.MagicMock(spec=Agent) + agent.instruction = "Initial instruction" + agent.sub_agents = {} + agent.clone.return_value = agent + return agent + + +def test_adapter_init(mock_gepa, mock_sampler, mock_agent): + del mock_gepa # only needed to mock gepa in background + loop = asyncio.new_event_loop() + _AdapterClass = _create_agent_gepa_adapter_class() + adapter = _AdapterClass(mock_agent, mock_sampler, loop) + assert adapter._initial_agent == mock_agent + assert adapter._sampler == mock_sampler + assert adapter._main_loop == loop + assert adapter._train_example_ids == {"train1", "train2"} + assert adapter._validation_example_ids == {"val1", "val2"} + loop.close() + + +def test_adapter_evaluate_train(mocker, mock_gepa, mock_sampler, mock_agent): + del mock_gepa # only needed to mock gepa in background + loop = mocker.MagicMock(spec=asyncio.AbstractEventLoop) + _AdapterClass = _create_agent_gepa_adapter_class() + adapter = _AdapterClass(mock_agent, mock_sampler, loop) + + candidate = {"agent_prompt": "New prompt"} + batch = ["train1"] + + # mock the future returned by run_coroutine_threadsafe + mock_future = mocker.MagicMock() + expected_result = UnstructuredSamplingResult( + scores={"train1": 0.8}, + data={"train1": {"output": "result"}}, + ) + mock_future.result.return_value = expected_result + + mock_rct = mocker.patch( + "asyncio.run_coroutine_threadsafe", return_value=mock_future + ) + eval_batch = adapter.evaluate(batch, candidate, capture_traces=True) + + mock_rct.assert_called_once() + mock_sampler.sample_and_score.assert_called_once_with( + mocker.ANY, + example_set="train", + batch=batch, + capture_full_eval_data=True, + ) + + mock_agent.clone.assert_called_once_with(update={"instruction": "New prompt"}) + + assert isinstance(eval_batch, MockEvaluationBatch) + assert eval_batch.scores == [0.8] + assert eval_batch.outputs == [{"output": "result"}] + assert eval_batch.trajectories == [{"output": "result"}] + + +def test_adapter_evaluate_validation( + mocker, mock_gepa, mock_sampler, mock_agent +): + del mock_gepa # only needed to mock gepa in background + loop = mocker.MagicMock(spec=asyncio.AbstractEventLoop) + _AdapterClass = _create_agent_gepa_adapter_class() + adapter = _AdapterClass(mock_agent, mock_sampler, loop) + + candidate = {"agent_prompt": "New prompt"} + batch = ["val1"] + + mock_future = mocker.MagicMock() + expected_result = UnstructuredSamplingResult(scores={"val1": 0.5}, data={}) + mock_future.result.return_value = expected_result + + mocker.patch("asyncio.run_coroutine_threadsafe", return_value=mock_future) + adapter.evaluate(batch, candidate) + + mock_sampler.sample_and_score.assert_called_once_with( + mocker.ANY, + example_set="validation", + batch=batch, + capture_full_eval_data=False, + ) + + +def test_adapter_make_reflective_dataset( + mocker, mock_gepa, mock_sampler, mock_agent +): + del mock_gepa # only needed to mock gepa in background + loop = mocker.MagicMock(spec=asyncio.AbstractEventLoop) + _AdapterClass = _create_agent_gepa_adapter_class() + adapter = _AdapterClass(mock_agent, mock_sampler, loop) + + candidate = {"agent_prompt": "Prompt"} + eval_batch = MockEvaluationBatch( + outputs=[{"o": 1}, {"o": 2}], + scores=[0.9, 0.1], + trajectories=[{"t": 1}, {"t": 2}], + ) + components = ["component1"] + + dataset = adapter.make_reflective_dataset(candidate, eval_batch, components) + + assert "component1" in dataset + assert len(dataset["component1"]) == 2 + assert dataset["component1"][0] == { + "agent_prompt": "Prompt", + "score": 0.9, + "eval_data": {"t": 1}, + } + assert dataset["component1"][1] == { + "agent_prompt": "Prompt", + "score": 0.1, + "eval_data": {"t": 2}, + } + + +@pytest.mark.asyncio +async def test_optimize(mocker, mock_gepa, mock_sampler, mock_agent): + config = GEPARootAgentPromptOptimizerConfig() + optimizer = GEPARootAgentPromptOptimizer(config) + + # mock LLM + mock_llm_class = mocker.MagicMock() + mock_llm = mocker.MagicMock() + mock_llm_class.return_value = mock_llm + optimizer._llm_class = mock_llm_class + + # mock gepa.optimize return value + mock_gepa_result = mocker.MagicMock() + mock_gepa_result.candidates = [{"agent_prompt": "Optimized instruction"}] + mock_gepa_result.val_aggregate_scores = [0.95] + mock_gepa_result.to_dict.return_value = {"full": "result"} + mock_gepa.optimize.return_value = mock_gepa_result + + result = await optimizer.optimize(mock_agent, mock_sampler) + + mock_gepa.optimize.assert_called_once() + call_kwargs = mock_gepa.optimize.call_args[1] + + assert call_kwargs["seed_candidate"] == { + "agent_prompt": "Initial instruction" + } + assert call_kwargs["trainset"] == ["train1", "train2"] + assert call_kwargs["valset"] == ["val1", "val2"] + + assert len(result.optimized_agents) == 1 + assert result.optimized_agents[0].overall_score == 0.95 + mock_agent.clone.assert_called_with( + update={"instruction": "Optimized instruction"} + ) + assert result.gepa_result == {"full": "result"} + + +@pytest.mark.asyncio +async def test_optimize_logs_warning_on_overlapping_ids( + mocker, mock_gepa, mock_sampler, mock_agent +): + # Setup overlapping IDs + mock_sampler.get_train_example_ids.return_value = ["id1", "id2"] + mock_sampler.get_validation_example_ids.return_value = ["id2", "id3"] + + config = GEPARootAgentPromptOptimizerConfig() + optimizer = GEPARootAgentPromptOptimizer(config) + + # Mock LLM class + mock_llm_class = mocker.MagicMock() + optimizer._llm_class = mock_llm_class + + # Mock gepa.optimize return value + mock_gepa_result = mocker.MagicMock() + mock_gepa_result.candidates = [] + mock_gepa_result.val_aggregate_scores = [] + mock_gepa_result.to_dict.return_value = {} + mock_gepa.optimize.return_value = mock_gepa_result + + mock_logger = mocker.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer._logger" + ) + + # Run optimization + await optimizer.optimize(mock_agent, mock_sampler) + + # Verify warning + mock_logger.warning.assert_called_with( + "The training and validation example UIDs overlap. This WILL cause" + " aliasing issues unless each common UID refers to the same example" + " in both sets." + )