feat: Improve gepa tau-bench colab for external use

PiperOrigin-RevId: 828579343

contributing/samples/gepa/adk_agent.py

@@ -263,30 +263,35 @@ def run_environment_loop(
           ),
       ],
   )
-  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
+
+  async def _async_run():
+    runner = runners.InMemoryRunner(
+        agent=agent,
+        app_name='eval_app',
+        plugins=plugins,
+    )
+    session = await 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.
+    events = [
+        event_lib.Event(
+            author='user',
+            content=initial_message,
+        )
+    ]
+    async for event in runner.run_async(
+        user_id=session.user_id,
+        session_id=session.id,
+        new_message=initial_message,
+    ):
+      events.append(event)
+    return events
+
+  return asyncio.run(_async_run())

contributing/samples/gepa/adk_agent_test.py

@@ -324,7 +324,10 @@ def test_model_name_is_set():
   async def _mock_create_session(*args, **kwargs):
     del args, kwargs
     await asyncio.sleep(0.1)
-    return
+    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
@@ -332,16 +335,14 @@ def test_model_name_is_set():
         _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([])],
-        )
+    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

contributing/samples/gepa/gepa_tau_bench.ipynb

@@ -31,11 +31,13 @@
       "execution_count": null,
       "metadata": {
         "id": "GqUHYdvRJ7pt",
-        "language": "python"
+        "language": "python",
+        "cellView": "form"
       },
       "outputs": [],
       "source": [
         "#@title Install Tau-bench and GEPA\n",
+        "!git clone https://github.com/google/adk-python.git\n",
         "!git clone https://github.com/sierra-research/tau-bench.git\n",
         "%cd tau-bench/\n",
         "!pip install -e . --quiet\n",
@@ -43,9 +45,39 @@
         "%cd ..\n",
         "!pip install gepa --quiet\n",
         "\n",
-        "!pip install retry --quiet"
+        "!pip install retry --quiet\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Configure python dependencies\n",
+        "import sys\n",
+        "\n",
+        "sys.path.append('/content/tau-bench')\n",
+        "sys.path.append('/content/adk-python/contributing/samples/gepa')"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "k0nrsIca0yXr"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Authentication\n",
+        "from google.colab import auth\n",
+        "auth.authenticate_user()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "NsXa217t03vL"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -55,7 +87,7 @@
       },
       "outputs": [],
       "source": [
-        "#@title Setup and Authentication\n",
+        "#@title Setup\n",
         "from datetime import datetime\n",
         "import json\n",
         "import logging\n",

contributing/samples/gepa/rater_lib.py

@@ -120,7 +120,7 @@ def format_user_agent_conversation(conv: list[dict[str, Any]]) -> str:
         res = (
             res
             + 'The execution result from the agent of function'
-            f' {fc["name"]} is: \n{fc["args"]}\n'
+            f' {fc["name"]} is: \n{fc["response"]}\n'
         )
   return res