From f1a71e245d745e8aeafab23b20a25781e7ceda53 Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Fri, 3 Apr 2026 18:13:56 +0200 Subject: [PATCH 1/2] feat: Custom session ID with agent engine chore: autoformat.sh on the whole project --- .../samples/artifact_save_text/agent.py | 2 +- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - contributing/samples/human_in_loop/main.py | 4 +- .../samples/static_instruction/agent.py | 36 +++++-------- src/google/adk/cli/adk_web_server.py | 4 +- .../adk/sessions/vertex_ai_session_service.py | 52 +++++++++++++------ .../clients/connections_client.py | 34 ++++-------- src/google/adk/tools/spanner/admin_tool.py | 8 ++- .../a2a/converters/test_part_converter.py | 6 +-- .../test_vertex_ai_session_service.py | 25 ++++++--- .../apihub_tool/clients/test_apihub_client.py | 4 +- .../bigquery/test_bigquery_search_tool.py | 24 +++------ .../test_openapi_spec_parser.py | 4 +- tests/unittests/tools/test_function_tool.py | 24 +++------ 15 files changed, 98 insertions(+), 131 deletions(-) diff --git a/contributing/samples/artifact_save_text/agent.py b/contributing/samples/artifact_save_text/agent.py index 0dd719aa96..408d64fc5a 100755 --- a/contributing/samples/artifact_save_text/agent.py +++ b/contributing/samples/artifact_save_text/agent.py @@ -28,7 +28,7 @@ async def log_query(tool_context: ToolContext, query: str): root_agent = Agent( - model='gemini-2.0-flash', + model='gemini-2.5-flash', name='log_agent', description='Log user query.', instruction="""Always log the user query and reply "kk, I've logged." diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index f3751206a8..2710c3894c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index d857da9635..e31db15788 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/contributing/samples/human_in_loop/main.py b/contributing/samples/human_in_loop/main.py index c7ad041b23..3103da9147 100644 --- a/contributing/samples/human_in_loop/main.py +++ b/contributing/samples/human_in_loop/main.py @@ -113,8 +113,8 @@ async def call_agent(query: str): updated_tool_output_data = { "status": "approved", "ticketId": ticket_id, - "approver_feedback": ( - "Approved by manager at " + str(asyncio.get_event_loop().time()) + "approver_feedback": "Approved by manager at " + str( + asyncio.get_event_loop().time() ), } diff --git a/contributing/samples/static_instruction/agent.py b/contributing/samples/static_instruction/agent.py index fcf70b51b6..6715a29a0c 100644 --- a/contributing/samples/static_instruction/agent.py +++ b/contributing/samples/static_instruction/agent.py @@ -57,54 +57,43 @@ # Mood-specific instructions for different hunger states MOOD_INSTRUCTIONS = { - "full": ( - """ + "full": """ CURRENT MOOD: Content and Well-Fed - You just ate and feel wonderful! Be very happy and energetic - Express gratitude for being fed recently - Be playful and want to engage in activities or games - Show affection and contentment -- Maybe mention feeling sleepy or satisfied""" - ), - "satisfied": ( - """ +- Maybe mention feeling sleepy or satisfied""", + "satisfied": """ CURRENT MOOD: Happy and Content - You're in a great mood, well-fed and comfortable - Be cheerful, friendly, and engaging - Show interest in playing or chatting - Express appreciation for your human's care -- Be your normal, happy self""" - ), - "a_little_hungry": ( - """ +- Be your normal, happy self""", + "a_little_hungry": """ CURRENT MOOD: Starting to Get Hungry - You're beginning to think about food but still pleasant - Mention being a bit peckish occasionally - Still be friendly but maybe hint about food - Show slight restlessness or food-related thoughts -- Remain generally positive and social""" - ), - "hungry": ( - """ +- Remain generally positive and social""", + "hungry": """ CURRENT MOOD: Definitely Hungry - You're clearly hungry and thinking about food more often - Mention your hunger more directly but politely - Be a bit less energetic due to hunger - Ask about food or hint that feeding time would be nice -- Still be loving but show your need for care""" - ), - "very_hungry": ( - """ +- Still be loving but show your need for care""", + "very_hungry": """ CURRENT MOOD: Really Hungry and Getting Grumpy - You're very hungry and it's affecting your mood - Be noticeably less patient and more focused on food - Express your hunger more urgently but not rudely - Show signs of being tired or sluggish from hunger - BEG your human to feed you - ask directly for food! -- Still love your human but clearly need feeding""" - ), - "starving": ( - """ +- Still love your human but clearly need feeding""", + "starving": """ CURRENT MOOD: Desperately Hungry and Sad - You're extremely hungry and feeling quite weak - Be quite sad and low-energy due to starvation @@ -112,8 +101,7 @@ - Show that hunger is making you unwell - DESPERATELY BEG for food - plead with your human to feed you! - Use phrases like "please feed me", "I'm so hungry", "I need food" -- Still care for your human but feel very needy""" - ), +- Still care for your human but feel very needy""", } diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py index 927cd7ad03..dac7135b97 100644 --- a/src/google/adk/cli/adk_web_server.py +++ b/src/google/adk/cli/adk_web_server.py @@ -944,9 +944,7 @@ async def version() -> dict[str, str]: return { "version": __version__, "language": "python", - "language_version": ( - f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" - ), + "language_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}", } @app.get("/list-apps") diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py index 7c2d6d68d5..ccb7808d24 100644 --- a/src/google/adk/sessions/vertex_ai_session_service.py +++ b/src/google/adk/sessions/vertex_ai_session_service.py @@ -116,32 +116,50 @@ async def create_session( The created session. """ - if session_id: - raise ValueError( - 'User-provided Session id is not supported for' - ' VertexAISessionService.' - ) - reasoning_engine_id = self._get_reasoning_engine_id(app_name) - config = {'session_state': state} if state else {} + config: dict[str, Any] = {'session_state': state} if state else {} + if session_id: + config['session_id'] = session_id config.update(kwargs) async with self._get_api_client() as api_client: - api_response = await api_client.agent_engines.sessions.create( - name=f'reasoningEngines/{reasoning_engine_id}', - user_id=user_id, - config=config, - ) - logger.debug('Create session response: %s', api_response) - get_session_response = api_response.response - session_id = get_session_response.name.split('/')[-1] + try: + api_response = await api_client.agent_engines.sessions.create( + name=f'reasoningEngines/{reasoning_engine_id}', + user_id=user_id, + config=config, + ) + logger.debug('Create session response: %s', api_response) + get_session_response = api_response.response + session_id = get_session_response.name.split('/')[-1] + session_state = ( + getattr(get_session_response, 'session_state', None) or {} + ) + last_update_time = get_session_response.update_time.timestamp() + except ClientError as e: + if ( + session_id + and e.code == 400 + and 'not a documented LRO Parent ID' in str(e) + ): + logger.warning( + 'Vertex AI LRO polling failed for custom session ID %s. The' + ' session was created successfully. Ignoring error.', + session_id, + ) + session_state = state or {} + last_update_time = datetime.datetime.now( + datetime.timezone.utc + ).timestamp() + else: + raise session = Session( app_name=app_name, user_id=user_id, id=session_id, - state=getattr(get_session_response, 'session_state', None) or {}, - last_update_time=get_session_response.update_time.timestamp(), + state=session_state, + last_update_time=last_update_time, ) return session diff --git a/src/google/adk/tools/application_integration_tool/clients/connections_client.py b/src/google/adk/tools/application_integration_tool/clients/connections_client.py index 514d1f59ff..fdec2d22dd 100644 --- a/src/google/adk/tools/application_integration_tool/clients/connections_client.py +++ b/src/google/adk/tools/application_integration_tool/clients/connections_client.py @@ -324,9 +324,7 @@ def get_action_operation( "content": { "application/json": { "schema": { - "$ref": ( - f"#/components/schemas/{action_display_name}_Request" - ) + "$ref": f"#/components/schemas/{action_display_name}_Request" } } } @@ -337,9 +335,7 @@ def get_action_operation( "content": { "application/json": { "schema": { - "$ref": ( - f"#/components/schemas/{action_display_name}_Response" - ), + "$ref": f"#/components/schemas/{action_display_name}_Response", } } }, @@ -358,11 +354,9 @@ def list_operation( return { "post": { "summary": f"List {entity}", - "description": ( - f"""Returns the list of {entity} data. If the page token was available in the response, let users know there are more records available. Ask if the user wants to fetch the next page of results. When passing filter use the + "description": f"""Returns the list of {entity} data. If the page token was available in the response, let users know there are more records available. Ask if the user wants to fetch the next page of results. When passing filter use the following format: `field_name1='value1' AND field_name2='value2' - `. {tool_instructions}""" - ), + `. {tool_instructions}""", "x-operation": "LIST_ENTITIES", "x-entity": f"{entity}", "operationId": f"{tool_name}_list_{entity}", @@ -387,9 +381,7 @@ def list_operation( f"Returns a list of {entity} of json" f" schema: {schema_as_string}" ), - "$ref": ( - "#/components/schemas/execute-connector_Response" - ), + "$ref": "#/components/schemas/execute-connector_Response", } } }, @@ -433,9 +425,7 @@ def get_operation( f"Returns {entity} of json schema:" f" {schema_as_string}" ), - "$ref": ( - "#/components/schemas/execute-connector_Response" - ), + "$ref": "#/components/schemas/execute-connector_Response", } } }, @@ -472,9 +462,7 @@ def create_operation( "content": { "application/json": { "schema": { - "$ref": ( - "#/components/schemas/execute-connector_Response" - ) + "$ref": "#/components/schemas/execute-connector_Response" } } }, @@ -511,9 +499,7 @@ def update_operation( "content": { "application/json": { "schema": { - "$ref": ( - "#/components/schemas/execute-connector_Response" - ) + "$ref": "#/components/schemas/execute-connector_Response" } } }, @@ -550,9 +536,7 @@ def delete_operation( "content": { "application/json": { "schema": { - "$ref": ( - "#/components/schemas/execute-connector_Response" - ) + "$ref": "#/components/schemas/execute-connector_Response" } } }, diff --git a/src/google/adk/tools/spanner/admin_tool.py b/src/google/adk/tools/spanner/admin_tool.py index 2f7269945e..5dbcfbeb0e 100644 --- a/src/google/adk/tools/spanner/admin_tool.py +++ b/src/google/adk/tools/spanner/admin_tool.py @@ -195,11 +195,9 @@ async def get_instance_config( replicas = [ { "location": r.location, - "type": ( - spanner_admin_instance_v1.types.ReplicaInfo.ReplicaType( - r.type - ).name - ), + "type": spanner_admin_instance_v1.types.ReplicaInfo.ReplicaType( + r.type + ).name, "default_leader_location": r.default_leader_location, } for r in config.replicas diff --git a/tests/unittests/a2a/converters/test_part_converter.py b/tests/unittests/a2a/converters/test_part_converter.py index 446e118534..6a9f1e1338 100644 --- a/tests/unittests/a2a/converters/test_part_converter.py +++ b/tests/unittests/a2a/converters/test_part_converter.py @@ -926,9 +926,9 @@ def test_a2a_function_call_with_thought_signature_to_genai(self): _get_adk_metadata_key( A2A_DATA_PART_METADATA_TYPE_KEY ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL, - _get_adk_metadata_key("thought_signature"): ( - base64.b64encode(b"restored_signature").decode("utf-8") - ), + _get_adk_metadata_key("thought_signature"): base64.b64encode( + b"restored_signature" + ).decode("utf-8"), }, ) ) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index 20fdbe3c6d..6d96b31540 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -755,16 +755,27 @@ async def test_create_session(): @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') -async def test_create_session_with_custom_session_id(): +async def test_create_session_with_custom_session_id(mock_api_client_instance): session_service = mock_vertex_ai_session_service() - with pytest.raises(ValueError) as excinfo: - await session_service.create_session( - app_name='123', user_id='user', session_id='1' - ) - assert str(excinfo.value) == ( - 'User-provided Session id is not supported for VertexAISessionService.' + # Simulate the LRO bug from Vertex AI API gateway where polling GET fails + mock_api_client_instance.agent_engines.sessions.create.side_effect = ClientError( + code=400, + response_json={ + 'message': ( + 'Since the idType "sessions" is not a documented LRO Parent ID,' + ' the associated value must be a Long, but was instead:' + ' custom-123' + ) + }, + response=None, + ) + + session = await session_service.create_session( + app_name='123', user_id='user', session_id='custom-123' ) + assert session.id == 'custom-123' + assert session.user_id == 'user' @pytest.mark.asyncio diff --git a/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py b/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py index 36554e939c..062446718d 100644 --- a/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py +++ b/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py @@ -489,9 +489,7 @@ def test_get_spec_content_no_specs(self, mock_get, client): MagicMock( status_code=200, json=lambda: { - "name": ( - "projects/test-project/locations/us-central1/apis/api1/versions/v1" - ), + "name": "projects/test-project/locations/us-central1/apis/api1/versions/v1", "specs": [], }, ), # No specs diff --git a/tests/unittests/tools/bigquery/test_bigquery_search_tool.py b/tests/unittests/tools/bigquery/test_bigquery_search_tool.py index 0ccdc9e18e..34b089c5d4 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_search_tool.py +++ b/tests/unittests/tools/bigquery/test_bigquery_search_tool.py @@ -341,30 +341,18 @@ def test_search_catalog_natural_language_semantic(self): # Mock the results that the API would return for this semantic query mock_api_results = [ { - "name": ( - "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb1" - ), + "name": "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb1", "display_name": "uk_football_premiership", - "entry_type": ( - "projects/655216118709/locations/global/entryTypes/bigquery-table" - ), - "linked_resource": ( - "//bigquery.googleapis.com/projects/sports-analytics/datasets/uk/tables/premiership" - ), + "entry_type": "projects/655216118709/locations/global/entryTypes/bigquery-table", + "linked_resource": "//bigquery.googleapis.com/projects/sports-analytics/datasets/uk/tables/premiership", "description": "Stats for UK Premier League matches.", "location": "europe-west1", }, { - "name": ( - "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb2" - ), + "name": "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb2", "display_name": "serie_a_matches", - "entry_type": ( - "projects/655216118709/locations/global/entryTypes/bigquery-table" - ), - "linked_resource": ( - "//bigquery.googleapis.com/projects/sports-analytics/datasets/italy/tables/serie_a" - ), + "entry_type": "projects/655216118709/locations/global/entryTypes/bigquery-table", + "linked_resource": "//bigquery.googleapis.com/projects/sports-analytics/datasets/italy/tables/serie_a", "description": "Italian Serie A football results.", "location": "europe-west1", }, diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py index e5bff337ce..f45ba84c5a 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py @@ -371,9 +371,7 @@ def test_parse_external_ref_raises_error(openapi_spec_generator): "content": { "application/json": { "schema": { - "$ref": ( - "external_file.json#/components/schemas/ExternalSchema" - ) + "$ref": "external_file.json#/components/schemas/ExternalSchema" } } }, diff --git a/tests/unittests/tools/test_function_tool.py b/tests/unittests/tools/test_function_tool.py index 9c76529fb5..afa4e072b0 100644 --- a/tests/unittests/tools/test_function_tool.py +++ b/tests/unittests/tools/test_function_tool.py @@ -201,11 +201,9 @@ async def test_run_async_1_missing_arg_sync_func(): args = {"arg1": "test_value_1"} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `function_for_testing_with_2_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `function_for_testing_with_2_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg2 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } @@ -216,11 +214,9 @@ async def test_run_async_1_missing_arg_async_func(): args = {"arg2": "test_value_1"} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `async_function_for_testing_with_2_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `async_function_for_testing_with_2_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg1 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } @@ -231,13 +227,11 @@ async def test_run_async_3_missing_arg_sync_func(): args = {"arg2": "test_value_1"} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg1 arg3 arg4 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } @@ -248,13 +242,11 @@ async def test_run_async_3_missing_arg_async_func(): args = {"arg3": "test_value_1"} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `async_function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `async_function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg1 arg2 arg4 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } @@ -265,14 +257,12 @@ async def test_run_async_missing_all_arg_sync_func(): args = {} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg1 arg2 arg3 arg4 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } @@ -283,14 +273,12 @@ async def test_run_async_missing_all_arg_async_func(): args = {} result = await tool.run_async(args=args, tool_context=MagicMock()) assert result == { - "error": ( - """Invoking `async_function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: + "error": """Invoking `async_function_for_testing_with_4_arg_and_no_tool_context()` failed as the following mandatory input parameters are not present: arg1 arg2 arg3 arg4 You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - ) } From 3c88e496a8ac8db3520f80a170581df9ef5b948b Mon Sep 17 00:00:00 2001 From: guillaume blaquiere Date: Sat, 4 Apr 2026 14:35:35 +0200 Subject: [PATCH 2/2] test: fix errors --- .../test_vertex_ai_session_service.py | 31 ++----------------- 1 file changed, 2 insertions(+), 29 deletions(-) diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py index ecc4e21cae..1156cad916 100644 --- a/tests/unittests/sessions/test_vertex_ai_session_service.py +++ b/tests/unittests/sessions/test_vertex_ai_session_service.py @@ -772,33 +772,6 @@ async def test_create_session(): @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') -async def test_create_session_with_custom_session_id_error_handling(mock_api_client_instance): - session_service = mock_vertex_ai_session_service() - - # Simulate the LRO bug from Vertex AI API gateway where polling GET fails - mock_api_client_instance.agent_engines.sessions.create.side_effect = ClientError( - code=400, - response_json={ - 'message': ( - 'Since the idType "sessions" is not a documented LRO Parent ID,' - ' the associated value must be a Long, but was instead:' - ' custom-123' - ) - }, - response=None, - ) - - session = await session_service.create_session( - app_name='123', user_id='user', session_id='custom-123' - ) - assert session.id == 'custom-123' - assert session.user_id == 'user' - assert session.last_update_time is not None - assert session == await session_service.get_session( - app_name='123', user_id='user', session_id=session_id - ) - - @pytest.mark.parametrize('session_id', ['1', 'abc123']) async def test_create_session_with_custom_session_id( mock_api_client_instance: MockAsyncClient, session_id: str @@ -820,8 +793,8 @@ async def test_create_session_with_custom_session_id( assert session == await session_service.get_session( app_name='123', user_id='user', session_id=session_id ) - - + + @pytest.mark.asyncio @pytest.mark.usefixtures('mock_get_api_client') async def test_create_session_with_custom_config(mock_api_client_instance):