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 <keyurj@google.com>
PiperOrigin-RevId: 878009649

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",

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(),
+    )

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."
+  )