feat: Add API key support for creating agent engines

PiperOrigin-RevId: 753758169

Author: vertex-sdk-bot

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
@@ -409,12 +409,22 @@ def create(
             )
             agent_engine_spec.class_methods.extend(class_methods_spec)
             reasoning_engine.spec = agent_engine_spec
-        operation_future = sdk_resource.api_client.create_reasoning_engine(
-            parent=initializer.global_config.common_location_path(
-                project=sdk_resource.project, location=sdk_resource.location
-            ),
-            reasoning_engine=reasoning_engine,
-        )
+        if (
+            initializer.global_config.project is not None
+            and initializer.global_config.location is not None
+        ):
+            operation_future = sdk_resource.api_client.create_reasoning_engine(
+                parent=initializer.global_config.common_location_path(
+                    project=sdk_resource.project, location=sdk_resource.location
+                ),
+                reasoning_engine=reasoning_engine,
+            )
+        elif initializer.global_config.api_key is not None:
+            operation_future = sdk_resource.api_client.create_reasoning_engine(
+                reasoning_engine=reasoning_engine,
+            )
+        else:
+            raise ValueError("Either project or api_key must be set in vertexai.init.")
         _LOGGER.log_create_with_lro(cls, operation_future)
         _LOGGER.info(
             f"View progress and logs at https://console.cloud.google.com/logs/query?project={sdk_resource.project}"
@@ -430,11 +440,21 @@ def create(
         sdk_resource._gca_resource = sdk_resource._get_gca_resource(
             resource_name=created_resource.name
         )
-        sdk_resource.execution_api_client = initializer.global_config.create_client(
-            client_class=aip_utils.AgentEngineExecutionClientWithOverride,
-            credentials=sdk_resource.credentials,
-            location_override=sdk_resource.location,
-        )
+        if (
+            initializer.global_config.project is not None
+            and initializer.global_config.location is not None
+        ):
+            sdk_resource.execution_api_client = initializer.global_config.create_client(
+                client_class=aip_utils.AgentEngineExecutionClientWithOverride,
+                credentials=sdk_resource.credentials,
+                location_override=sdk_resource.location,
+            )
+        elif initializer.global_config.api_key is not None:
+            sdk_resource.execution_api_client = initializer.global_config.create_client(
+                client_class=aip_utils.AgentEngineExecutionClientWithOverride,
+                api_key=initializer.global_config.api_key,
+                location_override=sdk_resource.location,
+            )
         if agent_engine is not None:
             try:
                 _register_api_methods_or_raise(sdk_resource)
@@ -881,17 +901,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 +1013,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()