chore: import upstream snapshot with attribution

tests/unittests/optimization/gepa_root_agent_optimizer_test.py

@@ -0,0 +1,452 @@
+# 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
+from collections.abc import Callable
+import sys
+from typing import Any
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.optimization import gepa_root_agent_optimizer
+from google.adk.optimization.data_types import UnstructuredSamplingResult
+from google.adk.optimization.gepa_root_agent_optimizer import _create_agent_from_candidate
+from google.adk.optimization.gepa_root_agent_optimizer import _create_agent_gepa_adapter_class
+from google.adk.optimization.gepa_root_agent_optimizer import _update_skill_toolset
+from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizer
+from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizerConfig
+from google.adk.optimization.sampler import Sampler
+from google.adk.skills import models
+from google.adk.tools.skill_toolset import SkillToolset
+import pytest
+
+# Spec structures used to autospec the dynamically mocked third-party `gepa`
+# package. Since gepa is a lazy-loaded dependency in the runtime code, it may
+# not be available in our standard hermetic test environment at import time.
+# These placeholders allow us to build strict type/interface checks using
+# `create_autospec` without requiring the gepa dependency.
+
+
+class MockEvaluationBatchSpec:
+
+  def __init__(self, outputs, scores, trajectories):
+    self.outputs = outputs
+    self.scores = scores
+    self.trajectories = trajectories
+
+
+class MockGEPAAdapterSpec:
+  """Mock that supports generic type hints."""
+
+  def __class_getitem__(cls, item):
+    return cls
+
+
+class MockAdapterModuleSpec:
+  EvaluationBatch = MockEvaluationBatchSpec
+  GEPAAdapter = MockGEPAAdapterSpec
+
+
+class MockInstructionProposalSignatureSpec:
+
+  @staticmethod
+  def prompt_renderer(input_dict):
+    pass
+
+  @staticmethod
+  def output_extractor(lm_out):
+    pass
+
+
+class MockInstructionProposalSpec:
+  InstructionProposalSignature = MockInstructionProposalSignatureSpec
+
+
+class MockStrategiesSpec:
+  instruction_proposal = MockInstructionProposalSpec
+
+
+class MockCoreSpec:
+  adapter_module = MockAdapterModuleSpec
+
+
+class MockGEPAModuleSpec:
+  core = MockCoreSpec
+  strategies = MockStrategiesSpec
+
+  @staticmethod
+  def optimize(*args, **kwargs):
+    pass
+
+
+class MockGEPAResultSpec:
+  candidates: list[dict[str, str]] = []
+  val_aggregate_scores: list[float] = []
+
+  def to_dict(self) -> dict[str, Any]:
+    return {}
+
+
+class MockSamplerSpec:
+
+  def get_train_example_ids(self) -> list[str]:
+    return []
+
+  def get_validation_example_ids(self) -> list[str]:
+    return []
+
+  def sample_and_score(self, *args, **kwargs):
+    pass
+
+
+@pytest.fixture(name="mock_gepa")
+def fixture_mock_gepa(mocker):
+  # mock gepa before it gets imported by the optimizer module
+  mock_gepa_module = mocker.create_autospec(MockGEPAModuleSpec)
+  mock_gepa_adapter_module = mocker.create_autospec(MockAdapterModuleSpec)
+
+  mock_gepa_adapter_module.EvaluationBatch = MockEvaluationBatchSpec
+  mock_gepa_adapter_module.GEPAAdapter = MockGEPAAdapterSpec
+
+  mock_gepa_module.core = mocker.create_autospec(MockCoreSpec)
+  mock_gepa_module.core.adapter = mock_gepa_adapter_module
+
+  mock_gepa_module.strategies = mocker.create_autospec(MockStrategiesSpec)
+  mock_ip = mocker.create_autospec(MockInstructionProposalSpec)
+  mock_gepa_module.strategies.instruction_proposal = mock_ip
+  mock_ip.InstructionProposalSignature = mocker.create_autospec(
+      MockInstructionProposalSignatureSpec
+  )
+
+  mocker.patch.dict(
+      sys.modules,
+      {
+          "gepa": mock_gepa_module,
+          "gepa.core": mock_gepa_module.core,
+          "gepa.core.adapter": mock_gepa_adapter_module,
+          "gepa.strategies": mock_gepa_module.strategies,
+          "gepa.strategies.instruction_proposal": (
+              mock_gepa_module.strategies.instruction_proposal
+          ),
+      },
+  )
+  return mock_gepa_module
+
+
+@pytest.fixture
+def mock_sampler(mocker):
+  sampler = mocker.create_autospec(MockSamplerSpec)
+  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.create_autospec(Agent, instance=True)
+  agent.instruction = "Initial instruction"
+  agent.sub_agents = {}
+  agent.clone.return_value = agent
+  agent.tools = []
+  return agent
+
+
+@pytest.fixture
+def mock_adapter(mocker, mock_gepa, mock_agent, mock_sampler):
+  del mock_gepa  # only needed to mock gepa in background
+  loop = mocker.create_autospec(asyncio.AbstractEventLoop, instance=True)
+  mock_reflection_lm = mocker.create_autospec(Callable)
+  _AdapterClass = _create_agent_gepa_adapter_class()
+  return _AdapterClass(mock_agent, mock_sampler, loop, mock_reflection_lm)
+
+
+def test_create_agent_from_candidate(mock_agent):
+  mock_agent.tools = []
+  candidate = {"agent_prompt": "New prompt"}
+  new_agent = _create_agent_from_candidate(mock_agent, candidate)
+
+  mock_agent.clone.assert_called_once_with(update={"instruction": "New prompt"})
+  assert new_agent == mock_agent
+
+
+def test_update_skill_toolset(mocker):
+  mock_skill = mocker.create_autospec(models.Skill, instance=True)
+  mock_skill.name = "my_skill"
+  mock_skill.instructions = "Old skill inst"
+  mock_skill_copy = mocker.create_autospec(models.Skill, instance=True)
+  mock_skill.model_copy.return_value = mock_skill_copy
+
+  mock_skill_toolset = mocker.create_autospec(SkillToolset, instance=True)
+  type(mock_skill_toolset).skills = mocker.PropertyMock(
+      return_value=[mock_skill]
+  )
+  mock_new_toolset = mocker.create_autospec(SkillToolset, instance=True)
+  mock_skill_toolset.clone_with_updated_skills.return_value = mock_new_toolset
+
+  candidate = {
+      "skill_instructions:my_skill": "New skill inst",
+  }
+
+  result = _update_skill_toolset(mock_skill_toolset, candidate)
+
+  mock_skill.model_copy.assert_called_once_with(
+      update={"instructions": "New skill inst"}
+  )
+  mock_skill_toolset.clone_with_updated_skills.assert_called_once_with(
+      [mock_skill_copy]
+  )
+  assert result is mock_new_toolset
+
+
+def test_create_agent_from_candidate_with_skills(mocker, mock_agent):
+  mock_skill_toolset = mocker.create_autospec(SkillToolset, instance=True)
+  mock_new_toolset = mocker.create_autospec(SkillToolset, instance=True)
+
+  mock_update = mocker.patch.object(
+      gepa_root_agent_optimizer,
+      "_update_skill_toolset",
+      return_value=mock_new_toolset,
+      autospec=True,
+  )
+
+  mock_agent.tools = [mock_skill_toolset]
+
+  candidate = {
+      "agent_prompt": "New prompt",
+      "skill_instructions:my_skill": "New skill inst",
+  }
+
+  new_agent = _create_agent_from_candidate(mock_agent, candidate)
+
+  mock_agent.clone.assert_called_once_with(update={"instruction": "New prompt"})
+  mock_update.assert_called_once_with(mock_skill_toolset, candidate)
+
+  assert len(new_agent.tools) == 1
+  assert new_agent.tools[0] is mock_new_toolset
+
+
+def test_adapter_init(mocker, 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()
+  mock_reflection_lm = mocker.create_autospec(Callable)
+  adapter = _AdapterClass(mock_agent, mock_sampler, loop, mock_reflection_lm)
+  assert adapter._initial_agent == mock_agent
+  assert adapter._sampler == mock_sampler
+  assert adapter._main_loop == loop
+  assert adapter._reflection_lm == mock_reflection_lm
+  assert adapter._train_example_ids == {"train1", "train2"}
+  assert adapter._validation_example_ids == {"val1", "val2"}
+  loop.close()
+
+
+def test_adapter_evaluate_train(mocker, mock_adapter, mock_sampler, mock_agent):
+  candidate = {"agent_prompt": "New prompt"}
+  batch = ["train1"]
+
+  # mock the future returned by run_coroutine_threadsafe
+  mock_future = mocker.create_autospec(asyncio.Future, instance=True)
+  expected_result = UnstructuredSamplingResult(
+      scores={"train1": 0.8},
+      data={"train1": {"output": "result"}},
+  )
+  mock_future.result.return_value = expected_result
+
+  mock_rct = mocker.patch.object(
+      asyncio,
+      "run_coroutine_threadsafe",
+      return_value=mock_future,
+      autospec=True,
+  )
+  eval_batch = mock_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, MockEvaluationBatchSpec)
+  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_adapter, mock_sampler):
+  candidate = {"agent_prompt": "New prompt"}
+  batch = ["val1"]
+
+  mock_future = mocker.create_autospec(asyncio.Future, instance=True)
+  expected_result = UnstructuredSamplingResult(scores={"val1": 0.5}, data={})
+  mock_future.result.return_value = expected_result
+
+  mocker.patch.object(
+      asyncio,
+      "run_coroutine_threadsafe",
+      return_value=mock_future,
+      autospec=True,
+  )
+  mock_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(mock_adapter):
+  candidate = {"agent_prompt": "Prompt"}
+  eval_batch = MockEvaluationBatchSpec(
+      outputs=[{"o": 1}, {"o": 2}],
+      scores=[0.9, 0.1],
+      trajectories=[{"t": "uses my_skill"}, {"t": "does not use skill"}],
+  )
+  components = ["agent_prompt", "skill_instructions:my_skill"]
+
+  dataset = mock_adapter.make_reflective_dataset(
+      candidate, eval_batch, components
+  )
+
+  assert dataset == {
+      "agent_prompt": [
+          {
+              "score": 0.9,
+              "eval_data": {"t": "uses my_skill"},
+          },
+          {
+              "score": 0.1,
+              "eval_data": {"t": "does not use skill"},
+          },
+      ],
+      "skill_instructions:my_skill": [
+          {
+              "score": 0.9,
+              "eval_data": {"t": "uses my_skill"},
+          },
+      ],
+  }
+
+
+def test_adapter_propose_new_texts(mock_gepa, mock_adapter):
+  mock_adapter._reflection_lm.return_value = "lm output"
+
+  candidate = {
+      "agent_prompt": "Old prompt",
+      "skill_instructions:my_skill": "Old skill inst",
+  }
+  reflective_dataset = {
+      "agent_prompt": [{"score": 1.0, "eval_data": {}}],
+      "skill_instructions:my_skill": [{"score": 0.9, "eval_data": {}}],
+  }
+  components = ["agent_prompt", "skill_instructions:my_skill"]
+
+  mock_ips = (
+      mock_gepa.strategies.instruction_proposal.InstructionProposalSignature
+  )
+  mock_ips.prompt_renderer.return_value = "rendered prompt"
+  mock_ips.output_extractor.side_effect = [
+      {"new_instruction": "New prompt"},
+      {"new_instruction": "New skill inst"},
+  ]
+
+  new_texts = mock_adapter.propose_new_texts(
+      candidate, reflective_dataset, components
+  )
+
+  assert mock_ips.prompt_renderer.call_count == 2
+  assert mock_adapter._reflection_lm.call_count == 2
+  assert mock_ips.output_extractor.call_count == 2
+  assert new_texts == {
+      "agent_prompt": "New prompt",
+      "skill_instructions:my_skill": "New skill inst",
+  }
+
+
+async def test_optimize(mocker, mock_gepa, mock_sampler, mock_agent):
+  config = GEPARootAgentOptimizerConfig()
+  optimizer = GEPARootAgentOptimizer(config)
+
+  # mock LLM
+  mock_llm_class = mocker.create_autospec(Callable)
+  mock_llm = mocker.create_autospec(Callable)
+  mock_llm_class.return_value = mock_llm
+  optimizer._llm_class = mock_llm_class
+
+  # mock gepa.optimize return value
+  mock_gepa_result = mocker.create_autospec(MockGEPAResultSpec, instance=True)
+  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"}
+
+
+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 = GEPARootAgentOptimizerConfig()
+  optimizer = GEPARootAgentOptimizer(config)
+
+  # Mock LLM class
+  mock_llm_class = mocker.create_autospec(Callable)
+  optimizer._llm_class = mock_llm_class
+
+  # Mock gepa.optimize return value
+  mock_gepa_result = mocker.create_autospec(MockGEPAResultSpec, instance=True)
+  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.object(
+      gepa_root_agent_optimizer, "logger", autospec=True
+  )
+
+  # 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."
+  )

tests/unittests/optimization/gepa_root_agent_prompt_optimizer_test.py

@@ -0,0 +1,265 @@
+# 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.mode = None
+  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."
+  )

tests/unittests/optimization/local_eval_sampler_test.py

@@ -0,0 +1,383 @@
+# 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
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.evaluation.base_eval_service import EvaluateConfig
+from google.adk.evaluation.base_eval_service import EvaluateRequest
+from google.adk.evaluation.base_eval_service import InferenceConfig
+from google.adk.evaluation.base_eval_service import InferenceRequest
+from google.adk.evaluation.base_eval_service import InferenceResult
+from google.adk.evaluation.eval_case import Invocation
+from google.adk.evaluation.eval_case import InvocationEvent
+from google.adk.evaluation.eval_case import InvocationEvents
+from google.adk.evaluation.eval_config import EvalConfig
+from google.adk.evaluation.eval_config import EvalMetric
+from google.adk.evaluation.eval_metrics import EvalMetricResult
+from google.adk.evaluation.eval_metrics import EvalMetricResultPerInvocation
+from google.adk.evaluation.eval_metrics import EvalStatus
+from google.adk.evaluation.eval_result import EvalCaseResult
+from google.adk.evaluation.eval_sets_manager import EvalSetsManager
+from google.adk.optimization.local_eval_sampler import _log_eval_summary
+from google.adk.optimization.local_eval_sampler import extract_single_invocation_info
+from google.adk.optimization.local_eval_sampler import extract_tool_call_data
+from google.adk.optimization.local_eval_sampler import LocalEvalSampler
+from google.adk.optimization.local_eval_sampler import LocalEvalSamplerConfig
+from google.genai import types
+import pytest
+
+
+def test_log_eval_summary(mocker):
+  statuses = (
+      [EvalStatus.PASSED] * 3
+      + [EvalStatus.FAILED] * 2
+      + [EvalStatus.NOT_EVALUATED]
+  )
+  expected_log = "Evaluation summary: 3 PASSED, 2 FAILED, 1 OTHER"
+
+  eval_results = [
+      mocker.MagicMock(spec=EvalCaseResult, final_eval_status=status)
+      for status in statuses
+  ]
+  mock_logger = mocker.patch(
+      "google.adk.optimization.local_eval_sampler.logger"
+  )
+
+  _log_eval_summary(eval_results)
+
+  mock_logger.info.assert_called_once_with(expected_log)
+
+
+def test_extract_tool_call_data():
+  # omitting IntermediateData tests as it is no longer used
+  # case 1: empty invocation events
+  assert not extract_tool_call_data(InvocationEvents())
+  # case 2: multi call invocation events
+  multi_call_invocation_events = InvocationEvents(
+      invocation_events=[
+          InvocationEvent(
+              author="agent",
+              content=types.Content(
+                  parts=[
+                      types.Part(
+                          function_call=types.FunctionCall(
+                              id="call_1",
+                              name="tool_1",
+                              args={"a": 1},
+                          )
+                      ),
+                      types.Part(
+                          function_call=types.FunctionCall(
+                              id="call_2",
+                              name="tool_2",
+                              args={"b": 2},
+                          )
+                      ),
+                      types.Part(
+                          function_response=types.FunctionResponse(
+                              id="call_1",
+                              name="tool_1",
+                              response={"result_1": "done"},
+                          )
+                      ),
+                      types.Part(
+                          function_response=types.FunctionResponse(
+                              id="call_2",
+                              name="tool_2",
+                              response={"result_2": "done"},
+                          )
+                      ),
+                  ]
+              ),
+          )
+      ]
+  )
+  expected_entries = [
+      {
+          "name": "tool_1",
+          "args": {"a": 1},
+          "response": {"result_1": "done"},
+      },
+      {
+          "name": "tool_2",
+          "args": {"b": 2},
+          "response": {"result_2": "done"},
+      },
+  ]
+  result = extract_tool_call_data(multi_call_invocation_events)
+  # order is not guaranteed
+  for expected_entry in expected_entries:
+    assert expected_entry in result
+  assert len(result) == len(expected_entries)
+
+
+def test_extract_single_invocation_info():
+  invocation = Invocation(
+      user_content=types.Content(
+          parts=[
+              types.Part(text="user thought", thought=True),
+              types.Part(text="Hello agent!"),
+          ]
+      ),
+      final_response=types.Content(
+          parts=[
+              types.Part(text="agent thought", thought=True),
+              types.Part(text="Hello user!"),
+          ]
+      ),
+  )
+
+  result = extract_single_invocation_info(invocation)
+
+  assert result == {
+      "user_prompt": "Hello agent!",
+      "agent_response": "Hello user!",
+  }
+
+
+@pytest.mark.parametrize(
+    "config_kwargs, expected_attrs",
+    [
+        (
+            {"train_eval_set": "train_set"},
+            {
+                "_train_eval_set": "train_set",
+                "_train_eval_case_ids": ["train_set_1", "train_set_2"],
+                "_validation_eval_set": "train_set",
+                "_validation_eval_case_ids": ["train_set_1", "train_set_2"],
+            },
+        ),
+        (
+            {"train_eval_set": "train_set", "train_eval_case_ids": ["t1"]},
+            {
+                "_train_eval_case_ids": ["t1"],
+                "_validation_eval_case_ids": ["t1"],
+            },
+        ),
+        (
+            {"train_eval_set": "train_set", "validation_eval_set": "val_set"},
+            {
+                "_validation_eval_set": "val_set",
+                "_validation_eval_case_ids": ["val_set_1", "val_set_2"],
+            },
+        ),
+        (
+            {"train_eval_set": "train_set", "validation_eval_case_ids": ["v1"]},
+            {
+                "_validation_eval_case_ids": ["v1"],
+            },
+        ),
+        (
+            {
+                "train_eval_set": "train_set",
+                "train_eval_case_ids": ["t1"],
+                "validation_eval_set": "val_set",
+                "validation_eval_case_ids": ["v1"],
+            },
+            {
+                "_train_eval_case_ids": ["t1"],
+                "_validation_eval_set": "val_set",
+                "_validation_eval_case_ids": ["v1"],
+            },
+        ),
+    ],
+)
+def test_local_eval_service_interface_init(
+    mocker, config_kwargs, expected_attrs
+):
+  mock_eval_sets_manager = mocker.MagicMock(spec=EvalSetsManager)
+
+  def mock_get_eval_case_ids(self, eval_set_id):
+    return [f"{eval_set_id}_1", f"{eval_set_id}_2"]
+
+  mocker.patch.object(
+      LocalEvalSampler,
+      "_get_eval_case_ids",
+      autospec=True,
+      side_effect=mock_get_eval_case_ids,
+  )
+
+  config = LocalEvalSamplerConfig(
+      eval_config=EvalConfig(), app_name="test_app", **config_kwargs
+  )
+  interface = LocalEvalSampler(config, mock_eval_sets_manager)
+
+  for attr, expected_value in expected_attrs.items():
+    assert getattr(interface, attr) == expected_value
+
+
+@pytest.mark.asyncio
+async def test_evaluate_agent(mocker):
+  # Mocking LocalEvalService and its methods
+  mock_eval_service_cls = mocker.patch(
+      "google.adk.optimization.local_eval_sampler.LocalEvalService"
+  )
+  mock_eval_service = mock_eval_service_cls.return_value
+
+  # mocking inference
+  mock_inference_result = mocker.MagicMock(spec=InferenceResult)
+
+  async def mock_perform_inference(*args, **kwargs):
+    yield mock_inference_result
+
+  mock_eval_service.perform_inference.side_effect = mock_perform_inference
+
+  # mocking evaluate
+  mock_eval_case_result = mocker.MagicMock(spec=EvalCaseResult)
+
+  async def mock_evaluate(*args, **kwargs):
+    yield mock_eval_case_result
+
+  mock_eval_service.evaluate.side_effect = mock_evaluate
+
+  # mocking get_eval_metrics_from_config
+  mock_metrics = [EvalMetric(metric_name="test_metric")]
+  mocker.patch(
+      "google.adk.optimization.local_eval_sampler.get_eval_metrics_from_config",
+      return_value=mock_metrics,
+  )
+
+  mocker.patch("google.adk.evaluation.base_eval_service.EvaluateConfig")
+
+  # Initialize Interface
+  config = LocalEvalSamplerConfig(
+      eval_config=EvalConfig(),
+      app_name="test_app",
+      train_eval_set="train_set",
+      train_eval_case_ids=["t1"],
+  )
+  interface = LocalEvalSampler(config, mocker.MagicMock(spec=EvalSetsManager))
+
+  # Call _evaluate_agent
+  results = await interface._evaluate_agent(
+      mocker.MagicMock(spec=Agent), "train_set", ["t1"]
+  )
+
+  # Assertions
+  mock_eval_service.perform_inference.assert_called_once_with(
+      inference_request=InferenceRequest(
+          app_name="test_app",
+          eval_set_id="train_set",
+          eval_case_ids=["t1"],
+          inference_config=InferenceConfig(),
+      )
+  )
+  mock_eval_service.evaluate.assert_called_once_with(
+      evaluate_request=EvaluateRequest(
+          inference_results=[mock_inference_result],
+          evaluate_config=EvaluateConfig(eval_metrics=mock_metrics),
+      )
+  )
+  assert results == [mock_eval_case_result]
+
+
+@pytest.mark.asyncio
+async def test_extract_eval_data(mocker):
+  # Mock components
+  mock_eval_sets_manager = mocker.MagicMock(spec=EvalSetsManager)
+  mock_eval_case = mocker.MagicMock()
+  mock_eval_case.conversation_scenario = "test_scenario"
+  mock_eval_sets_manager.get_eval_case.return_value = mock_eval_case
+
+  # Mock per invocation result
+  mock_actual_invocation = mocker.MagicMock(spec=Invocation)
+  mock_expected_invocation = mocker.MagicMock(spec=Invocation)
+  mock_metric_result = mocker.MagicMock(spec=EvalMetricResult)
+  mock_metric_result.metric_name = "test_metric"
+  mock_metric_result.score = 0.854  # should be rounded to 0.85
+  mock_metric_result.eval_status = EvalStatus.PASSED
+
+  mock_per_inv_result = mocker.MagicMock(spec=EvalMetricResultPerInvocation)
+  mock_per_inv_result.actual_invocation = mock_actual_invocation
+  mock_per_inv_result.expected_invocation = mock_expected_invocation
+  mock_per_inv_result.eval_metric_results = [mock_metric_result]
+
+  mock_eval_result = mocker.MagicMock(spec=EvalCaseResult)
+  mock_eval_result.eval_id = "t1"
+  mock_eval_result.eval_metric_result_per_invocation = [mock_per_inv_result]
+
+  # Mock extract_single_invocation_info
+  mocker.patch(
+      "google.adk.optimization.local_eval_sampler.extract_single_invocation_info",
+      side_effect=[{"info": "actual"}, {"info": "expected"}],
+  )
+
+  # Initialize Interface
+  config = LocalEvalSamplerConfig(
+      eval_config=EvalConfig(),
+      app_name="test_app",
+      train_eval_set="train_set",
+      train_eval_case_ids=["t1"],
+  )
+  interface = LocalEvalSampler(config, mock_eval_sets_manager)
+
+  # Call _extract_eval_data
+  eval_data = interface._extract_eval_data("train_set", [mock_eval_result])
+
+  # Assertions
+  assert "t1" in eval_data
+  assert eval_data["t1"]["conversation_scenario"] == "test_scenario"
+  assert len(eval_data["t1"]["invocations"]) == 1
+  inv = eval_data["t1"]["invocations"][0]
+  assert inv["actual_invocation"] == {"info": "actual"}
+  assert inv["expected_invocation"] == {"info": "expected"}
+  assert inv["eval_metric_results"] == [
+      {"metric_name": "test_metric", "score": 0.85, "eval_status": "PASSED"}
+  ]
+
+
+@pytest.mark.asyncio
+async def test_sample_and_score(mocker):
+  # Mock results
+  mock_eval_result_1 = mocker.MagicMock(spec=EvalCaseResult)
+  mock_eval_result_1.eval_id = "t1"
+  mock_eval_result_1.final_eval_status = EvalStatus.PASSED
+
+  mock_eval_result_2 = mocker.MagicMock(spec=EvalCaseResult)
+  mock_eval_result_2.eval_id = "t2"
+  mock_eval_result_2.final_eval_status = EvalStatus.FAILED
+
+  eval_results = [mock_eval_result_1, mock_eval_result_2]
+
+  # Initialize Interface
+  config = LocalEvalSamplerConfig(
+      eval_config=EvalConfig(),
+      app_name="test_app",
+      train_eval_set="train_set",
+      train_eval_case_ids=["t1", "t2"],
+  )
+  interface = LocalEvalSampler(config, mocker.MagicMock(spec=EvalSetsManager))
+
+  # Patch internal methods
+  mocker.patch.object(interface, "_evaluate_agent", return_value=eval_results)
+  mock_log_summary = mocker.patch(
+      "google.adk.optimization.local_eval_sampler._log_eval_summary"
+  )
+  mock_extract_data = mocker.patch.object(
+      interface, "_extract_eval_data", return_value={"t1": {}, "t2": {}}
+  )
+
+  # Call sample_and_score
+  result = await interface.sample_and_score(
+      mocker.MagicMock(spec=Agent),
+      example_set="train",
+      capture_full_eval_data=True,
+  )
+
+  # Assertions
+  assert result.scores == {"t1": 1.0, "t2": 0.0}
+  assert result.data == {"t1": {}, "t2": {}}
+  mock_log_summary.assert_called_once_with(eval_results)
+  mock_extract_data.assert_called_once_with("train_set", eval_results)

tests/unittests/optimization/simple_prompt_optimizer_test.py

@@ -0,0 +1,104 @@
+# 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.
+
+"""Tests for simple_prompt_optimizer."""
+
+from unittest import mock
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.google_llm import LlmResponse
+from google.adk.optimization.data_types import UnstructuredSamplingResult
+from google.adk.optimization.sampler import Sampler
+from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer
+from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig
+from google.genai import types as genai_types
+import pytest
+
+
+@pytest.fixture
+def mock_sampler() -> mock.MagicMock:
+  sampler = mock.MagicMock(spec=Sampler)
+  sampler.get_train_example_ids.return_value = ["1", "2", "3", "4", "5"]
+  sampler.get_validation_example_ids.return_value = ["v1", "v2"]
+
+  async def mock_sample_and_score(
+      agent: Agent,
+      example_set: str,
+      batch: list[str] | None = None,
+      capture_full_eval_data: bool = False,
+  ) -> UnstructuredSamplingResult:
+    # Determine the actual batch to use
+    if batch is None:
+      if example_set == "train":
+        current_batch = sampler.get_train_example_ids()
+      else:  # "validation"
+        current_batch = sampler.get_validation_example_ids()
+    else:
+      current_batch = batch
+
+    # Simulate better score for "improved" prompt
+    if "IMPROVED" in agent.instruction:
+      scores = {uid: 0.9 for uid in current_batch}
+    else:
+      scores = {uid: 0.5 for uid in current_batch}
+    return UnstructuredSamplingResult(scores=scores)
+
+  sampler.sample_and_score.side_effect = mock_sample_and_score
+  return sampler
+
+
+@pytest.fixture
+def mock_llm_class() -> mock.MagicMock:
+  mock_llm = mock.MagicMock()
+
+  async def mock_generate_content_async(*args, **kwargs):
+    yield LlmResponse(
+        content=genai_types.Content(
+            parts=[genai_types.Part(text="IMPROVED PROMPT")]
+        )
+    )
+
+  mock_llm.generate_content_async.side_effect = mock_generate_content_async
+  mock_class = mock.MagicMock(return_value=mock_llm)
+  return mock_class
+
+
+@mock.patch(
+    "google.adk.optimization.simple_prompt_optimizer.LLMRegistry.resolve"
+)
+@pytest.mark.asyncio
+async def test_simple_prompt_optimizer(
+    mock_llm_resolve: mock.MagicMock,
+    mock_llm_class: mock.MagicMock,
+    mock_sampler: mock.MagicMock,
+):
+  """Test the SimplePromptOptimizer."""
+  mock_llm_resolve.return_value = mock_llm_class
+  config = SimplePromptOptimizerConfig(num_iterations=2, batch_size=2)
+  optimizer = SimplePromptOptimizer(config)
+
+  initial_agent = Agent(name="test_agent", instruction="Initial Prompt")
+  result = await optimizer.optimize(initial_agent, mock_sampler)
+
+  # Assertions
+  assert len(result.optimized_agents) == 1
+  optimized_agent = result.optimized_agents[0].optimized_agent
+  assert optimized_agent.instruction == "IMPROVED PROMPT"
+  assert result.optimized_agents[0].overall_score == 0.9
+
+  # Check mock calls
+  assert mock_sampler.get_train_example_ids.call_count == 1
+  # 1 initial, 2 iterations, 1 final validation
+  assert mock_sampler.sample_and_score.call_count == 4
+  assert mock_llm_class.return_value.generate_content_async.call_count == 2