feat: Remove the requirement for a staging bucket to be set if an agent engine is created without extra packages or requirements files.

vertex-sdk-bot · copybara-github · commit 6de9de172548 · 2025-05-02T15:11:19.000-07:00
PiperOrigin-RevId: 754147393
diff --git a/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py
@@ -2315,6 +2315,47 @@ def test_invalid_operation_schema(
         assert want_log_output in caplog.text
 
 
+@pytest.mark.usefixtures("google_auth_mock")
+class TestLightweightAgentEngine:
+    def setup_method(self):
+        importlib.reload(initializer)
+        importlib.reload(aiplatform)
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            credentials=_TEST_CREDENTIALS,
+        )
+        self.test_agent = CapitalizeEngine()
+
+    def test_create_agent_engine_with_no_spec(
+        self,
+        create_agent_engine_mock,
+        cloud_storage_create_bucket_mock,
+        tarfile_open_mock,
+        cloudpickle_dump_mock,
+        cloudpickle_load_mock,
+        importlib_metadata_version_mock,
+        get_agent_engine_mock,
+    ):
+        importlib.reload(initializer)
+        importlib.reload(aiplatform)
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            credentials=_TEST_CREDENTIALS,
+        )
+        agent_engines.create(
+            display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+            description=_TEST_AGENT_ENGINE_DESCRIPTION,
+        )
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            credentials=_TEST_CREDENTIALS,
+            staging_bucket=_TEST_STAGING_BUCKET,
+        )
+
+
 def _generate_agent_engine_to_update() -> "agent_engines.AgentEngine":
     test_agent_engine = agent_engines.create(CapitalizeEngine())
     # Resource name is required for the update method.
diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py
@@ -338,8 +338,11 @@ def create(
         """
         sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
         _validate_sys_version_or_raise(sys_version)
+        gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME
+        staging_bucket = initializer.global_config.staging_bucket
         if agent_engine is not None:
             agent_engine = _validate_agent_engine_or_raise(agent_engine)
+            _validate_staging_bucket_or_raise(staging_bucket)
         if agent_engine is None:
             if requirements is not None:
                 raise ValueError("requirements must be None if agent_engine is None.")
@@ -350,12 +353,9 @@ def create(
             requirements=requirements,
         )
         extra_packages = _validate_extra_packages_or_raise(extra_packages)
-        gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME
 
         sdk_resource = cls.__new__(cls)
         base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource)
-        staging_bucket = initializer.global_config.staging_bucket
-        _validate_staging_bucket_or_raise(staging_bucket)
         # Prepares the Agent Engine for creation in Vertex AI.
         # This involves packaging and uploading the artifacts for
         # agent_engine, requirements and extra_packages to
@@ -881,17 +881,18 @@ def _prepare(
         gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to
             use for staging the artifacts needed.
     """
+    if agent_engine is None:
+        return
     gcs_bucket = _get_gcs_bucket(
         project=project,
         location=location,
         staging_bucket=staging_bucket,
     )
-    if agent_engine is not None:
-        _upload_agent_engine(
-            agent_engine=agent_engine,
-            gcs_bucket=gcs_bucket,
-            gcs_dir_name=gcs_dir_name,
-        )
+    _upload_agent_engine(
+        agent_engine=agent_engine,
+        gcs_bucket=gcs_bucket,
+        gcs_dir_name=gcs_dir_name,
+    )
     if requirements is not None:
         _upload_requirements(
             requirements=requirements,
@@ -992,7 +993,7 @@ def _generate_update_request_or_raise(
         Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
     ] = None,
 ) -> reasoning_engine_service.UpdateReasoningEngineRequest:
-    """Tries to generates the update request for the agent engine."""
+    """Tries to generate the update request for the agent engine."""
     is_spec_update = False
     update_masks: List[str] = []
     agent_engine_spec = aip_types.ReasoningEngineSpec()
