chore: import upstream snapshot with attribution

contributing/samples/integrations/gepa/adk_agent_test.py

@@ -0,0 +1,349 @@
+# 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.
+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-plugin")
+    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)
+    mock_session = mock.Mock()
+    mock.user_id = "fake-user=id"
+    mock.id = "fake-session-id"
+    return mock_session
+
+  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 = []
+    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.model == "some-test-model"