Skip to content

Commit fb5ec5f

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
chore: Support detailed log filter in Agent Engine
PiperOrigin-RevId: 839479774
1 parent 2c362b3 commit fb5ec5f

File tree

3 files changed

+224
-8
lines changed

3 files changed

+224
-8
lines changed

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 176 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1521,6 +1521,89 @@ def test_prepare_without_creds(
15211521
gcs_dir_name=_TEST_GCS_DIR_NAME,
15221522
)
15231523

1524+
@pytest.mark.parametrize(
1525+
"operation_name,resource_name,expected_id,expected_exception,expected_message",
1526+
[
1527+
(
1528+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456",
1529+
"",
1530+
_TEST_RESOURCE_ID,
1531+
None,
1532+
None,
1533+
),
1534+
(
1535+
"",
1536+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1537+
_TEST_RESOURCE_ID,
1538+
None,
1539+
None,
1540+
),
1541+
(
1542+
"projects/123/locations/us-central1/reasoningEngines/other_id/operations/456",
1543+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1544+
_TEST_RESOURCE_ID,
1545+
None,
1546+
None,
1547+
),
1548+
(
1549+
"",
1550+
"",
1551+
None,
1552+
ValueError,
1553+
"Resource name or operation name cannot be empty.",
1554+
),
1555+
(
1556+
"invalid/operation/name",
1557+
"",
1558+
None,
1559+
ValueError,
1560+
"Failed to parse reasoning engine ID from operation name",
1561+
),
1562+
(
1563+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1564+
"",
1565+
None,
1566+
ValueError,
1567+
"Failed to parse reasoning engine ID from operation name",
1568+
),
1569+
(
1570+
"",
1571+
"invalid/resource/name",
1572+
None,
1573+
ValueError,
1574+
"Failed to parse reasoning engine ID from resource name",
1575+
),
1576+
(
1577+
"",
1578+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456",
1579+
None,
1580+
ValueError,
1581+
"Failed to parse reasoning engine ID from resource name",
1582+
),
1583+
],
1584+
)
1585+
def test_get_reasoning_engine_id(
1586+
self,
1587+
operation_name,
1588+
resource_name,
1589+
expected_id,
1590+
expected_exception,
1591+
expected_message,
1592+
):
1593+
if expected_exception:
1594+
with pytest.raises(expected_exception) as excinfo:
1595+
_agent_engines_utils._get_reasoning_engine_id(
1596+
operation_name=operation_name, resource_name=resource_name
1597+
)
1598+
assert expected_message in str(excinfo.value)
1599+
else:
1600+
assert (
1601+
_agent_engines_utils._get_reasoning_engine_id(
1602+
operation_name=operation_name, resource_name=resource_name
1603+
)
1604+
== expected_id
1605+
)
1606+
15241607

15251608
@pytest.mark.usefixtures("google_auth_mock")
15261609
class TestAgentEngine:
@@ -1570,7 +1653,18 @@ def test_list_agent_engine(self):
15701653
@pytest.mark.usefixtures("caplog")
15711654
@mock.patch.object(_agent_engines_utils, "_prepare")
15721655
@mock.patch.object(_agent_engines_utils, "_await_operation")
1573-
def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
1656+
@mock.patch.object(
1657+
_agent_engines_utils,
1658+
"_get_reasoning_engine_id",
1659+
return_value=_TEST_RESOURCE_ID,
1660+
)
1661+
def test_create_agent_engine(
1662+
self,
1663+
mock_get_reasoning_engine_id,
1664+
mock_await_operation,
1665+
mock_prepare,
1666+
caplog,
1667+
):
15741668
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
15751669
response=_genai_types.ReasoningEngine(
15761670
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
@@ -1620,8 +1714,14 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
16201714

16211715
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
16221716
@mock.patch.object(_agent_engines_utils, "_await_operation")
1717+
@mock.patch.object(
1718+
_agent_engines_utils,
1719+
"_get_reasoning_engine_id",
1720+
return_value=_TEST_RESOURCE_ID,
1721+
)
16231722
def test_create_agent_engine_lightweight(
16241723
self,
1724+
mock_get_reasoning_engine_id,
16251725
mock_await_operation,
16261726
mock_create_config,
16271727
):
@@ -1657,8 +1757,14 @@ def test_create_agent_engine_lightweight(
16571757

16581758
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
16591759
@mock.patch.object(_agent_engines_utils, "_await_operation")
1760+
@mock.patch.object(
1761+
_agent_engines_utils,
1762+
"_get_reasoning_engine_id",
1763+
return_value=_TEST_RESOURCE_ID,
1764+
)
16601765
def test_create_agent_engine_with_env_vars_dict(
16611766
self,
1767+
mock_get_reasoning_engine_id,
16621768
mock_await_operation,
16631769
mock_create_config,
16641770
):
@@ -1748,8 +1854,14 @@ def test_create_agent_engine_with_env_vars_dict(
17481854

17491855
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
17501856
@mock.patch.object(_agent_engines_utils, "_await_operation")
1857+
@mock.patch.object(
1858+
_agent_engines_utils,
1859+
"_get_reasoning_engine_id",
1860+
return_value=_TEST_RESOURCE_ID,
1861+
)
17511862
def test_create_agent_engine_with_custom_service_account(
17521863
self,
1864+
mock_get_reasoning_engine_id,
17531865
mock_await_operation,
17541866
mock_create_config,
17551867
):
@@ -1842,8 +1954,14 @@ def test_create_agent_engine_with_custom_service_account(
18421954

18431955
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
18441956
@mock.patch.object(_agent_engines_utils, "_await_operation")
1957+
@mock.patch.object(
1958+
_agent_engines_utils,
1959+
"_get_reasoning_engine_id",
1960+
return_value=_TEST_RESOURCE_ID,
1961+
)
18451962
def test_create_agent_engine_with_experimental_mode(
18461963
self,
1964+
mock_get_reasoning_engine_id,
18471965
mock_await_operation,
18481966
mock_create_config,
18491967
):
@@ -1939,8 +2057,14 @@ def test_create_agent_engine_with_experimental_mode(
19392057
return_value="test_tarball",
19402058
)
19412059
@mock.patch.object(_agent_engines_utils, "_await_operation")
2060+
@mock.patch.object(
2061+
_agent_engines_utils,
2062+
"_get_reasoning_engine_id",
2063+
return_value=_TEST_RESOURCE_ID,
2064+
)
19422065
def test_create_agent_engine_with_source_packages(
19432066
self,
2067+
mock_get_reasoning_engine_id,
19442068
mock_await_operation,
19452069
mock_create_base64_encoded_tarball,
19462070
):
@@ -2001,8 +2125,14 @@ def test_create_agent_engine_with_source_packages(
20012125

20022126
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
20032127
@mock.patch.object(_agent_engines_utils, "_await_operation")
2128+
@mock.patch.object(
2129+
_agent_engines_utils,
2130+
"_get_reasoning_engine_id",
2131+
return_value=_TEST_RESOURCE_ID,
2132+
)
20042133
def test_create_agent_engine_with_class_methods(
20052134
self,
2135+
mock_get_reasoning_engine_id,
20062136
mock_await_operation,
20072137
mock_create_config,
20082138
):
@@ -2088,8 +2218,14 @@ def test_create_agent_engine_with_class_methods(
20882218

20892219
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
20902220
@mock.patch.object(_agent_engines_utils, "_await_operation")
2221+
@mock.patch.object(
2222+
_agent_engines_utils,
2223+
"_get_reasoning_engine_id",
2224+
return_value=_TEST_RESOURCE_ID,
2225+
)
20912226
def test_create_agent_engine_with_agent_framework(
20922227
self,
2228+
mock_get_reasoning_engine_id,
20932229
mock_await_operation,
20942230
mock_create_config,
20952231
):
@@ -2696,8 +2832,17 @@ def test_operation_schemas(
26962832
@mock.patch.object(_agent_engines_utils, "_prepare")
26972833
@mock.patch.object(agent_engines.AgentEngines, "_create")
26982834
@mock.patch.object(_agent_engines_utils, "_await_operation")
2835+
@mock.patch.object(
2836+
_agent_engines_utils,
2837+
"_get_reasoning_engine_id",
2838+
return_value=_TEST_RESOURCE_ID,
2839+
)
26992840
def test_create_agent_engine_with_creds(
2700-
self, mock_await_operation, mock_create, mock_prepare
2841+
self,
2842+
mock_get_reasoning_engine_id,
2843+
mock_await_operation,
2844+
mock_create,
2845+
mock_prepare,
27012846
):
27022847
mock_operation = mock.Mock()
27032848
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2728,8 +2873,18 @@ def test_create_agent_engine_with_creds(
27282873
@mock.patch.object(agent_engines.AgentEngines, "_create")
27292874
@mock.patch("google.auth.default")
27302875
@mock.patch.object(_agent_engines_utils, "_await_operation")
2876+
@mock.patch.object(
2877+
_agent_engines_utils,
2878+
"_get_reasoning_engine_id",
2879+
return_value=_TEST_RESOURCE_ID,
2880+
)
27312881
def test_create_agent_engine_without_creds(
2732-
self, mock_await_operation, mock_auth_default, mock_create, mock_prepare
2882+
self,
2883+
mock_get_reasoning_engine_id,
2884+
mock_await_operation,
2885+
mock_auth_default,
2886+
mock_create,
2887+
mock_prepare,
27332888
):
27342889
mock_operation = mock.Mock()
27352890
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2765,8 +2920,17 @@ def test_create_agent_engine_without_creds(
27652920
@mock.patch.object(_agent_engines_utils, "_prepare")
27662921
@mock.patch.object(agent_engines.AgentEngines, "_create")
27672922
@mock.patch.object(_agent_engines_utils, "_await_operation")
2923+
@mock.patch.object(
2924+
_agent_engines_utils,
2925+
"_get_reasoning_engine_id",
2926+
return_value=_TEST_RESOURCE_ID,
2927+
)
27682928
def test_create_agent_engine_with_no_creds_in_client(
2769-
self, mock_await_operation, mock_create, mock_prepare
2929+
self,
2930+
mock_get_reasoning_engine_id,
2931+
mock_await_operation,
2932+
mock_create,
2933+
mock_prepare,
27702934
):
27712935
mock_operation = mock.Mock()
27722936
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2812,7 +2976,14 @@ def setup_method(self):
28122976

28132977
@mock.patch.object(_agent_engines_utils, "_prepare")
28142978
@mock.patch.object(_agent_engines_utils, "_await_operation")
2815-
def test_create_agent_engine_error(self, mock_await_operation, mock_prepare):
2979+
@mock.patch.object(
2980+
_agent_engines_utils,
2981+
"_get_reasoning_engine_id",
2982+
return_value=_TEST_RESOURCE_ID,
2983+
)
2984+
def test_create_agent_engine_error(
2985+
self, mock_get_reasoning_engine_id, mock_await_operation, mock_prepare
2986+
):
28162987
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
28172988
error=_TEST_AGENT_ENGINE_ERROR,
28182989
)

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import json
2424
import logging
2525
import os
26+
import re
2627
import sys
2728
import tarfile
2829
import time
@@ -414,6 +415,39 @@ async def __call__(
414415
pass
415416

416417

418+
def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "") -> str:
419+
"""Returns reasoning engine ID from operation name or resource name."""
420+
if not resource_name and not operation_name:
421+
raise ValueError("Resource name or operation name cannot be empty.")
422+
423+
if resource_name:
424+
match = re.match(
425+
r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)$",
426+
resource_name,
427+
)
428+
if match:
429+
return match.group(1)
430+
else:
431+
raise ValueError(
432+
"Failed to parse reasoning engine ID from resource name: "
433+
f"`{resource_name}`"
434+
)
435+
436+
if not operation_name:
437+
raise ValueError("Operation name cannot be empty.")
438+
439+
match = re.match(
440+
r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)/operations/[^/]+$",
441+
operation_name,
442+
)
443+
if match:
444+
return match.group(1)
445+
raise ValueError(
446+
"Failed to parse reasoning engine ID from operation name: "
447+
f"`{operation_name}`"
448+
)
449+
450+
417451
async def _await_async_operation(
418452
*,
419453
operation_name: str,

‎vertexai/_genai/agent_engines.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -941,9 +941,14 @@ def create(
941941
build_options=config.build_options,
942942
)
943943
operation = self._create(config=api_config)
944-
# TODO: Use a more specific link.
944+
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
945+
operation_name=operation.name
946+
)
945947
logger.info(
946-
f"View progress and logs at https://console.cloud.google.com/logs/query?project={self._api_client.project}."
948+
"View progress and logs at https://console.cloud.google.com/logs/query?"
949+
f"project={self._api_client.project}"
950+
"&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A"
951+
f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22."
947952
)
948953
if agent is not None or config.source_packages is not None:
949954
poll_interval_seconds = 10
@@ -1476,8 +1481,14 @@ def update(
14761481
build_options=config.build_options,
14771482
)
14781483
operation = self._update(name=name, config=api_config)
1484+
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
1485+
resource_name=name
1486+
)
14791487
logger.info(
1480-
f"View progress and logs at https://console.cloud.google.com/logs/query?project={self._api_client.project}."
1488+
"View progress and logs at https://console.cloud.google.com/logs/query?"
1489+
f"project={self._api_client.project}"
1490+
"&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A"
1491+
f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22."
14811492
)
14821493
operation = _agent_engines_utils._await_operation(
14831494
operation_name=operation.name,

0 commit comments

Comments
 (0)