feat: add logging for agent engine creation

PiperOrigin-RevId: 780862041

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -15,6 +15,7 @@
 import asyncio
 import importlib
 import json
+import logging
 import os
 import pytest
 import sys
@@ -963,10 +964,17 @@ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
+ caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines")
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1004,6 +1012,12 @@ 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
+ )
 
  @mock.patch.object(agent_engines.AgentEngines, "_create_config")
  @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
@@ -1016,7 +1030,12 @@ def test_create_agent_engine_lightweight(
  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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1061,7 +1080,12 @@ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1106,10 +1130,19 @@ def test_create_agent_engine_with_env_vars_dict(
  None,
  )
 
+ @pytest.mark.usefixtures("caplog")
  @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()
+ def test_update_agent_engine_requirements(
+ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
+ caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines")
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1149,13 +1182,23 @@ def test_update_agent_engine_requirements(self, mock_await_operation, mock_prepa
  },
  None,
  )
+ assert "Agent Engine updated. 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_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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1201,8 +1244,15 @@ 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()
+ def test_update_agent_engine_env_vars(
+ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1251,7 +1301,12 @@ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:
@@ -1275,7 +1330,12 @@ 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,
+ spec=_TEST_AGENT_ENGINE_SPEC,
+ )
+ )
  with mock.patch.object(
  self.client.agent_engines._api_client, "request"
  ) as request_mock:

vertexai/_genai/agent_engines.py

@@ -1988,6 +1988,11 @@ def create(
  env_vars=config.env_vars,
  )
  operation = self._create(config=api_config)
+ # TODO: Use a more specific link.
+ 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:
@@ -2002,6 +2007,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(
+ f"agent_engine=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.
@@ -2043,11 +2052,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
@@ -2063,6 +2075,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")
@@ -2099,6 +2112,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
@@ -2305,12 +2319,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(
+ f"agent_engine=client.agent_engines.get('{agent.api_resource.name}')"
+ )
  return self._register_api_methods(agent=agent)
 
  def _stream_query(