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