From 3eb38bf1f566545f5a6aba16b356cb2e4c148992 Mon Sep 17 00:00:00 2001 From: Shawn Yang Date: Thu, 4 Dec 2025 14:37:49 -0800 Subject: [PATCH] chore: Support detailed log filter in Agent Engine PiperOrigin-RevId: 840407609 --- .../unit/vertexai/genai/test_agent_engines.py | 181 +++++++++++++++++- vertexai/_genai/_agent_engines_utils.py | 34 ++++ vertexai/_genai/agent_engines.py | 17 +- 3 files changed, 224 insertions(+), 8 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index a9560ebaa8..d51edce6ec 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -1521,6 +1521,89 @@ def test_prepare_without_creds( gcs_dir_name=_TEST_GCS_DIR_NAME, ) + @pytest.mark.parametrize( + "operation_name,resource_name,expected_id,expected_exception,expected_message", + [ + ( + f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456", + "", + _TEST_RESOURCE_ID, + None, + None, + ), + ( + "", + f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}", + _TEST_RESOURCE_ID, + None, + None, + ), + ( + "projects/123/locations/us-central1/reasoningEngines/other_id/operations/456", + f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}", + _TEST_RESOURCE_ID, + None, + None, + ), + ( + "", + "", + None, + ValueError, + "Resource name or operation name cannot be empty.", + ), + ( + "invalid/operation/name", + "", + None, + ValueError, + "Failed to parse reasoning engine ID from operation name", + ), + ( + f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}", + "", + None, + ValueError, + "Failed to parse reasoning engine ID from operation name", + ), + ( + "", + "invalid/resource/name", + None, + ValueError, + "Failed to parse reasoning engine ID from resource name", + ), + ( + "", + f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456", + None, + ValueError, + "Failed to parse reasoning engine ID from resource name", + ), + ], + ) + def test_get_reasoning_engine_id( + self, + operation_name, + resource_name, + expected_id, + expected_exception, + expected_message, + ): + if expected_exception: + with pytest.raises(expected_exception) as excinfo: + _agent_engines_utils._get_reasoning_engine_id( + operation_name=operation_name, resource_name=resource_name + ) + assert expected_message in str(excinfo.value) + else: + assert ( + _agent_engines_utils._get_reasoning_engine_id( + operation_name=operation_name, resource_name=resource_name + ) + == expected_id + ) + @pytest.mark.usefixtures("google_auth_mock") class TestAgentEngine: @@ -1570,7 +1653,18 @@ def test_list_agent_engine(self): @pytest.mark.usefixtures("caplog") @mock.patch.object(_agent_engines_utils, "_prepare") @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog): + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_agent_engine( + self, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_prepare, + caplog, + ): mock_await_operation.return_value = _genai_types.AgentEngineOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, @@ -1620,8 +1714,14 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog): @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_lightweight( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -1657,8 +1757,14 @@ def test_create_agent_engine_lightweight( @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_env_vars_dict( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -1748,8 +1854,14 @@ def test_create_agent_engine_with_env_vars_dict( @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_custom_service_account( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -1842,8 +1954,14 @@ def test_create_agent_engine_with_custom_service_account( @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_experimental_mode( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -1939,8 +2057,14 @@ def test_create_agent_engine_with_experimental_mode( return_value="test_tarball", ) @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_source_packages( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_base64_encoded_tarball, ): @@ -2001,8 +2125,14 @@ def test_create_agent_engine_with_source_packages( @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_class_methods( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -2088,8 +2218,14 @@ def test_create_agent_engine_with_class_methods( @mock.patch.object(agent_engines.AgentEngines, "_create_config") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_agent_framework( self, + mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): @@ -2696,8 +2832,17 @@ def test_operation_schemas( @mock.patch.object(_agent_engines_utils, "_prepare") @mock.patch.object(agent_engines.AgentEngines, "_create") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_creds( - self, mock_await_operation, mock_create, mock_prepare + self, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_create, + mock_prepare, ): mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME @@ -2728,8 +2873,18 @@ def test_create_agent_engine_with_creds( @mock.patch.object(agent_engines.AgentEngines, "_create") @mock.patch("google.auth.default") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_without_creds( - self, mock_await_operation, mock_auth_default, mock_create, mock_prepare + self, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_auth_default, + mock_create, + mock_prepare, ): mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME @@ -2765,8 +2920,17 @@ def test_create_agent_engine_without_creds( @mock.patch.object(_agent_engines_utils, "_prepare") @mock.patch.object(agent_engines.AgentEngines, "_create") @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) def test_create_agent_engine_with_no_creds_in_client( - self, mock_await_operation, mock_create, mock_prepare + self, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_create, + mock_prepare, ): mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME @@ -2812,7 +2976,14 @@ def setup_method(self): @mock.patch.object(_agent_engines_utils, "_prepare") @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_create_agent_engine_error(self, mock_await_operation, mock_prepare): + @mock.patch.object( + _agent_engines_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_agent_engine_error( + self, mock_get_reasoning_engine_id, mock_await_operation, mock_prepare + ): mock_await_operation.return_value = _genai_types.AgentEngineOperation( error=_TEST_AGENT_ENGINE_ERROR, ) diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index 987277ca47..f4cf78425d 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -23,6 +23,7 @@ import json import logging import os +import re import sys import tarfile import time @@ -414,6 +415,39 @@ async def __call__( pass +def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "") -> str: + """Returns reasoning engine ID from operation name or resource name.""" + if not resource_name and not operation_name: + raise ValueError("Resource name or operation name cannot be empty.") + + if resource_name: + match = re.match( + r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)$", + resource_name, + ) + if match: + return match.group(1) + else: + raise ValueError( + "Failed to parse reasoning engine ID from resource name: " + f"`{resource_name}`" + ) + + if not operation_name: + raise ValueError("Operation name cannot be empty.") + + match = re.match( + r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)/operations/[^/]+$", + operation_name, + ) + if match: + return match.group(1) + raise ValueError( + "Failed to parse reasoning engine ID from operation name: " + f"`{operation_name}`" + ) + + async def _await_async_operation( *, operation_name: str, diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 6cfcb8be62..258ab4c516 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -941,9 +941,14 @@ def create( build_options=config.build_options, ) operation = self._create(config=api_config) - # TODO: Use a more specific link. + reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( + operation_name=operation.name + ) logger.info( - f"View progress and logs at https://console.cloud.google.com/logs/query?project={self._api_client.project}." + "View progress and logs at https://console.cloud.google.com/logs/query?" + f"project={self._api_client.project}" + "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" + f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." ) if agent is not None or config.source_packages is not None: poll_interval_seconds = 10 @@ -1476,8 +1481,14 @@ def update( build_options=config.build_options, ) operation = self._update(name=name, config=api_config) + reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( + resource_name=name + ) logger.info( - f"View progress and logs at https://console.cloud.google.com/logs/query?project={self._api_client.project}." + "View progress and logs at https://console.cloud.google.com/logs/query?" + f"project={self._api_client.project}" + "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" + f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." ) operation = _agent_engines_utils._await_operation( operation_name=operation.name,