Skip to content

Commit 795ee17

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
feat: add logging for agent engine creation
PiperOrigin-RevId: 780862041
1 parent 2332c28 commit 795ee17

File tree

4 files changed

+136
-31
lines changed

4 files changed

+136
-31
lines changed

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

Lines changed: 71 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import asyncio
1616
import importlib
1717
import json
18+
import logging
1819
import os
1920
import pytest
2021
import sys
@@ -963,10 +964,17 @@ def test_list_agent_engine(self):
963964
None,
964965
)
965966

967+
@pytest.mark.usefixtures("caplog")
966968
@mock.patch.object(_agent_engines, "_prepare")
967969
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
968-
def test_create_agent_engine(self, mock_await_operation, mock_prepare):
969-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
970+
def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
971+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
972+
response=_genai_types.ReasoningEngine(
973+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
974+
spec=_TEST_AGENT_ENGINE_SPEC,
975+
)
976+
)
977+
caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines")
970978
with mock.patch.object(
971979
self.client.agent_engines._api_client, "request"
972980
) as request_mock:
@@ -1004,6 +1012,12 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare):
10041012
},
10051013
None,
10061014
)
1015+
assert "View progress and logs at" in caplog.text
1016+
assert "Agent Engine created. To use it in another session:" in caplog.text
1017+
assert (
1018+
f"agent_engine=client.agent_engines.get("
1019+
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
1020+
)
10071021

10081022
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
10091023
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
@@ -1016,7 +1030,12 @@ def test_create_agent_engine_lightweight(
10161030
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
10171031
description=_TEST_AGENT_ENGINE_DESCRIPTION,
10181032
)
1019-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1033+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1034+
response=_genai_types.ReasoningEngine(
1035+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1036+
spec=_TEST_AGENT_ENGINE_SPEC,
1037+
)
1038+
)
10201039
with mock.patch.object(
10211040
self.client.agent_engines._api_client, "request"
10221041
) as request_mock:
@@ -1061,7 +1080,12 @@ def test_create_agent_engine_with_env_vars_dict(
10611080
"agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK,
10621081
},
10631082
}
1064-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1083+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1084+
response=_genai_types.ReasoningEngine(
1085+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1086+
spec=_TEST_AGENT_ENGINE_SPEC,
1087+
)
1088+
)
10651089
with mock.patch.object(
10661090
self.client.agent_engines._api_client, "request"
10671091
) as request_mock:
@@ -1106,10 +1130,19 @@ def test_create_agent_engine_with_env_vars_dict(
11061130
None,
11071131
)
11081132

1133+
@pytest.mark.usefixtures("caplog")
11091134
@mock.patch.object(_agent_engines, "_prepare")
11101135
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
1111-
def test_update_agent_engine_requirements(self, mock_await_operation, mock_prepare):
1112-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1136+
def test_update_agent_engine_requirements(
1137+
self, mock_await_operation, mock_prepare, caplog
1138+
):
1139+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1140+
response=_genai_types.ReasoningEngine(
1141+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1142+
spec=_TEST_AGENT_ENGINE_SPEC,
1143+
)
1144+
)
1145+
caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines")
11131146
with mock.patch.object(
11141147
self.client.agent_engines._api_client, "request"
11151148
) as request_mock:
@@ -1149,13 +1182,23 @@ def test_update_agent_engine_requirements(self, mock_await_operation, mock_prepa
11491182
},
11501183
None,
11511184
)
1185+
assert "Agent Engine updated. To use it in another session:" in caplog.text
1186+
assert (
1187+
f"agent_engine=client.agent_engines.get("
1188+
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
1189+
)
11521190

11531191
@mock.patch.object(_agent_engines, "_prepare")
11541192
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
11551193
def test_update_agent_engine_extra_packages(
11561194
self, mock_await_operation, mock_prepare
11571195
):
1158-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1196+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1197+
response=_genai_types.ReasoningEngine(
1198+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1199+
spec=_TEST_AGENT_ENGINE_SPEC,
1200+
)
1201+
)
11591202
with mock.patch.object(
11601203
self.client.agent_engines._api_client, "request"
11611204
) as request_mock:
@@ -1201,8 +1244,15 @@ def test_update_agent_engine_extra_packages(
12011244

12021245
@mock.patch.object(_agent_engines, "_prepare")
12031246
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
1204-
def test_update_agent_engine_env_vars(self, mock_await_operation, mock_prepare):
1205-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1247+
def test_update_agent_engine_env_vars(
1248+
self, mock_await_operation, mock_prepare, caplog
1249+
):
1250+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1251+
response=_genai_types.ReasoningEngine(
1252+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1253+
spec=_TEST_AGENT_ENGINE_SPEC,
1254+
)
1255+
)
12061256
with mock.patch.object(
12071257
self.client.agent_engines._api_client, "request"
12081258
) as request_mock:
@@ -1251,7 +1301,12 @@ def test_update_agent_engine_env_vars(self, mock_await_operation, mock_prepare):
12511301

12521302
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
12531303
def test_update_agent_engine_display_name(self, mock_await_operation):
1254-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1304+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1305+
response=_genai_types.ReasoningEngine(
1306+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1307+
spec=_TEST_AGENT_ENGINE_SPEC,
1308+
)
1309+
)
12551310
with mock.patch.object(
12561311
self.client.agent_engines._api_client, "request"
12571312
) as request_mock:
@@ -1275,7 +1330,12 @@ def test_update_agent_engine_display_name(self, mock_await_operation):
12751330

12761331
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
12771332
def test_update_agent_engine_description(self, mock_await_operation):
1278-
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
1333+
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
1334+
response=_genai_types.ReasoningEngine(
1335+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1336+
spec=_TEST_AGENT_ENGINE_SPEC,
1337+
)
1338+
)
12791339
with mock.patch.object(
12801340
self.client.agent_engines._api_client, "request"
12811341
) as request_mock:

‎vertexai/_genai/agent_engines.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1988,6 +1988,11 @@ def create(
19881988
env_vars=config.env_vars,
19891989
)
19901990
operation = self._create(config=api_config)
1991+
# TODO: Use a more specific link.
1992+
logger.info(
1993+
"View progress and logs at"
1994+
f" https://console.cloud.google.com/logs/query?project={self._api_client.project}."
1995+
)
19911996
if agent_engine is None:
19921997
poll_interval_seconds = 1 # Lightweight agent engine resource creation.
19931998
else:
@@ -2002,6 +2007,10 @@ def create(
20022007
api_async_client=AsyncAgentEngines(api_client_=self._api_client),
20032008
api_resource=operation.response,
20042009
)
2010+
logger.info("Agent Engine created. To use it in another session:")
2011+
logger.info(
2012+
f"agent_engine=client.agent_engines.get('{agent.api_resource.name}')"
2013+
)
20052014
if agent_engine is not None:
20062015
# If the user did not provide an agent_engine (e.g. lightweight
20072016
# provisioning), it will not have any API methods registered.
@@ -2043,11 +2052,14 @@ def _create_config(
20432052
if agent_engine is not None:
20442053
sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
20452054
gcs_dir_name = gcs_dir_name or _agent_engines._DEFAULT_GCS_DIR_NAME
2046-
agent_engine = _agent_engines._validate_agent_engine_or_raise(agent_engine)
2055+
agent_engine = _agent_engines._validate_agent_engine_or_raise(
2056+
agent_engine=agent_engine, logger=logger
2057+
)
20472058
_agent_engines._validate_staging_bucket_or_raise(staging_bucket)
20482059
requirements = _agent_engines._validate_requirements_or_raise(
20492060
agent_engine=agent_engine,
20502061
requirements=requirements,
2062+
logger=logger,
20512063
)
20522064
extra_packages = _agent_engines._validate_extra_packages_or_raise(
20532065
extra_packages
@@ -2063,6 +2075,7 @@ def _create_config(
20632075
staging_bucket=staging_bucket,
20642076
gcs_dir_name=gcs_dir_name,
20652077
extra_packages=extra_packages,
2078+
logger=logger,
20662079
)
20672080
# Update the package spec.
20682081
update_masks.append("spec.package_spec.pickle_object_gcs_uri")
@@ -2099,6 +2112,7 @@ def _create_config(
20992112
class_methods = _agent_engines._generate_class_methods_spec_or_raise(
21002113
agent_engine=agent_engine,
21012114
operations=_agent_engines._get_registered_operations(agent_engine),
2115+
logger=logger,
21022116
)
21032117
agent_engine_spec["class_methods"] = [
21042118
_utils.to_dict(class_method) for class_method in class_methods
@@ -2305,12 +2319,20 @@ def update(
23052319
env_vars=config.env_vars,
23062320
)
23072321
operation = self._update(name=name, config=api_config)
2322+
logger.info(
2323+
"View progress and logs at"
2324+
f" https://console.cloud.google.com/logs/query?project={self._api_client.project}."
2325+
)
23082326
operation = self._await_operation(operation_name=operation.name)
23092327
agent = types.AgentEngine(
23102328
api_client=self,
23112329
api_async_client=AsyncAgentEngines(api_client_=self._api_client),
23122330
api_resource=operation.response,
23132331
)
2332+
logger.info("Agent Engine updated. To use it in another session:")
2333+
logger.info(
2334+
f"agent_engine=client.agent_engines.get('{agent.api_resource.name}')"
2335+
)
23142336
return self._register_api_methods(agent=agent)
23152337

23162338
def _stream_query(

0 commit comments

Comments
 (0)