From e02f177790d9772dd253c9102b80df1a9418aa7f Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 5 Nov 2025 12:23:20 -0800 Subject: [PATCH] feat: Improve gepa tau-bench colab for external use PiperOrigin-RevId: 828579343 --- contributing/samples/gepa/adk_agent.py | 59 ++++++++++--------- contributing/samples/gepa/adk_agent_test.py | 23 ++++---- .../samples/gepa/gepa_tau_bench.ipynb | 38 +++++++++++- contributing/samples/gepa/rater_lib.py | 2 +- 4 files changed, 80 insertions(+), 42 deletions(-) diff --git a/contributing/samples/gepa/adk_agent.py b/contributing/samples/gepa/adk_agent.py index 15e153da..e4bc517d 100644 --- a/contributing/samples/gepa/adk_agent.py +++ b/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()) diff --git a/contributing/samples/gepa/adk_agent_test.py b/contributing/samples/gepa/adk_agent_test.py index 2eea7325..7a7a98e2 100644 --- a/contributing/samples/gepa/adk_agent_test.py +++ b/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 diff --git a/contributing/samples/gepa/gepa_tau_bench.ipynb b/contributing/samples/gepa/gepa_tau_bench.ipynb index 448fdc8d..975a4631 100644 --- a/contributing/samples/gepa/gepa_tau_bench.ipynb +++ b/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", diff --git a/contributing/samples/gepa/rater_lib.py b/contributing/samples/gepa/rater_lib.py index 739bba07..332bf333 100644 --- a/contributing/samples/gepa/rater_lib.py +++ b/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