Skip to content

feat: add logging for agent engine creation #5494

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 66 additions & 27 deletions tests/unit/vertexai/genai/test_agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -521,16 +521,16 @@ def register_operations(self) -> Dict[str, List[str]]:
OperationRegistrableEngine().custom_method,
schema_name=_TEST_CUSTOM_METHOD_NAME,
)
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_STANDARD_API_MODE
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_STANDARD_API_MODE
)
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _utils.generate_schema(
OperationRegistrableEngine().custom_async_method,
schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
)
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_ASYNC_API_MODE
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_ASYNC_API_MODE
)
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _utils.generate_schema(
StreamQueryEngine().stream_query,
schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME,
Expand All @@ -540,23 +540,23 @@ def register_operations(self) -> Dict[str, List[str]]:
OperationRegistrableEngine().custom_stream_method,
schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME,
)
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_STREAM_API_MODE
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_STREAM_API_MODE
)
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _utils.generate_schema(
AsyncStreamQueryEngine().async_stream_query,
schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME,
)
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_ASYNC_STREAM_API_MODE
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_ASYNC_STREAM_API_MODE
)
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = _utils.generate_schema(
OperationRegistrableEngine().custom_async_stream_method,
schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
)
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_ASYNC_STREAM_API_MODE
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_ASYNC_STREAM_API_MODE
)
_TEST_OPERATION_REGISTRABLE_SCHEMAS = [
_TEST_AGENT_ENGINE_QUERY_SCHEMA,
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA,
Expand All @@ -581,9 +581,9 @@ def register_operations(self) -> Dict[str, List[str]]:
MethodToBeUnregisteredEngine().method_to_be_unregistered,
schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME,
)
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[
_TEST_MODE_KEY_IN_SCHEMA
] = _TEST_STANDARD_API_MODE
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
_TEST_STANDARD_API_MODE
)
_TEST_ASYNC_QUERY_SCHEMAS = [_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA]
_TEST_STREAM_QUERY_SCHEMAS = [
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA,
Expand Down Expand Up @@ -963,10 +963,13 @@ def test_list_agent_engine(self):
None,
)

@pytest.mark.usefixtures("caplog")
@mock.patch.object(_agent_engines, "_prepare")
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_create_agent_engine(self, mock_await_operation, mock_prepare):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down Expand Up @@ -1004,19 +1007,29 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare):
},
None,
)
assert "View progress and logs at" in caplog.text
assert "Agent Engine created. To use it in another session:" in caplog.text
assert (
f"agent_engine = client.agent_engines.get("
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
)

@pytest.mark.usefixtures("caplog")
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_create_agent_engine_lightweight(
self,
mock_await_operation,
mock_create_config,
caplog,
):
mock_create_config.return_value = _genai_types.CreateAgentEngineConfig(
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
description=_TEST_AGENT_ENGINE_DESCRIPTION,
)
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand All @@ -1040,13 +1053,21 @@ def test_create_agent_engine_lightweight(
},
None,
)
assert "View progress and logs at" in caplog.text
assert "Agent Engine created. To use it in another session:" in caplog.text
assert (
f"agent_engine = client.agent_engines.get("
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
)

@pytest.mark.usefixtures("caplog")
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_create_agent_engine_with_env_vars_dict(
self,
mock_await_operation,
mock_create_config,
caplog,
):
mock_create_config.return_value = {
"display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME,
Expand All @@ -1061,7 +1082,9 @@ def test_create_agent_engine_with_env_vars_dict(
"agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK,
},
}
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down Expand Up @@ -1105,11 +1128,19 @@ def test_create_agent_engine_with_env_vars_dict(
},
None,
)
assert "View progress and logs at" in caplog.text
assert "Agent Engine created. To use it in another session:" in caplog.text
assert (
f"agent_engine = client.agent_engines.get("
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
)

@mock.patch.object(_agent_engines, "_prepare")
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_update_agent_engine_requirements(self, mock_await_operation, mock_prepare):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down Expand Up @@ -1155,7 +1186,9 @@ def test_update_agent_engine_requirements(self, mock_await_operation, mock_prepa
def test_update_agent_engine_extra_packages(
self, mock_await_operation, mock_prepare
):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down Expand Up @@ -1202,7 +1235,9 @@ def test_update_agent_engine_extra_packages(
@mock.patch.object(_agent_engines, "_prepare")
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_update_agent_engine_env_vars(self, mock_await_operation, mock_prepare):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down Expand Up @@ -1251,7 +1286,9 @@ def test_update_agent_engine_env_vars(self, mock_await_operation, mock_prepare):

@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_update_agent_engine_display_name(self, mock_await_operation):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand All @@ -1275,7 +1312,9 @@ def test_update_agent_engine_display_name(self, mock_await_operation):

@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
def test_update_agent_engine_description(self, mock_await_operation):
mock_await_operation.return_value = _genai_types.AgentEngineOperation()
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
response=_genai_types.ReasoningEngine(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
)
with mock.patch.object(
self.client.agent_engines._api_client, "request"
) as request_mock:
Expand Down
23 changes: 22 additions & 1 deletion vertexai/_genai/agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -1988,6 +1988,10 @@ def create(
env_vars=config.env_vars,
)
operation = self._create(config=api_config)
logger.info(
"View progress and logs at"
f" https://console.cloud.google.com/logs/query?project={self._api_client.project}."
)
if agent_engine is None:
poll_interval_seconds = 1 # Lightweight agent engine resource creation.
else:
Expand All @@ -2002,6 +2006,10 @@ def create(
api_async_client=AsyncAgentEngines(api_client_=self._api_client),
api_resource=operation.response,
)
logger.info("Agent Engine created. To use it in another session:")
logger.info(
"agent_engine =" f" client.agent_engines.get('{agent.api_resource.name}')"
)
if agent_engine is not None:
# If the user did not provide an agent_engine (e.g. lightweight
# provisioning), it will not have any API methods registered.
Expand Down Expand Up @@ -2043,11 +2051,14 @@ def _create_config(
if agent_engine is not None:
sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
gcs_dir_name = gcs_dir_name or _agent_engines._DEFAULT_GCS_DIR_NAME
agent_engine = _agent_engines._validate_agent_engine_or_raise(agent_engine)
agent_engine = _agent_engines._validate_agent_engine_or_raise(
agent_engine=agent_engine, logger=logger
)
_agent_engines._validate_staging_bucket_or_raise(staging_bucket)
requirements = _agent_engines._validate_requirements_or_raise(
agent_engine=agent_engine,
requirements=requirements,
logger=logger,
)
extra_packages = _agent_engines._validate_extra_packages_or_raise(
extra_packages
Expand All @@ -2063,6 +2074,7 @@ def _create_config(
staging_bucket=staging_bucket,
gcs_dir_name=gcs_dir_name,
extra_packages=extra_packages,
logger=logger,
)
# Update the package spec.
update_masks.append("spec.package_spec.pickle_object_gcs_uri")
Expand Down Expand Up @@ -2099,6 +2111,7 @@ def _create_config(
class_methods = _agent_engines._generate_class_methods_spec_or_raise(
agent_engine=agent_engine,
operations=_agent_engines._get_registered_operations(agent_engine),
logger=logger,
)
agent_engine_spec["class_methods"] = [
_utils.to_dict(class_method) for class_method in class_methods
Expand Down Expand Up @@ -2305,12 +2318,20 @@ def update(
env_vars=config.env_vars,
)
operation = self._update(name=name, config=api_config)
logger.info(
"View progress and logs at"
f" https://console.cloud.google.com/logs/query?project={self._api_client.project}."
)
operation = self._await_operation(operation_name=operation.name)
agent = types.AgentEngine(
api_client=self,
api_async_client=AsyncAgentEngines(api_client_=self._api_client),
api_resource=operation.response,
)
logger.info("Agent Engine updated. To use it in another session:")
logger.info(
"agent_engine =" f" client.agent_engines.get('{agent.api_resource.name}')"
)
return self._register_api_methods(agent=agent)

def _stream_query(
Expand Down
Loading
Loading