feat: In run_query_job, rename gcs_bucket to gcs_uri and allow the case that user sets the filename for the output.

PiperOrigin-RevId: 893135854

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_agent_engines.py

@@ -2948,7 +2948,7 @@ def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_clie
                 name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
                 config={
                     "query": _TEST_QUERY_PROMPT,
-                    "gcs_bucket": "gs://my-input-bucket/",
+                    "output_gcs_uri": "gs://my-input-bucket/",
                 },
             )
 
@@ -2959,17 +2959,17 @@ def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_clie
 
             assert result == _genai_types.RunQueryJobResult(
                 job_name="projects/123/locations/us-central1/reasoningEngines/456/operations/789",
-                input_gcs_uri="gs://my-input-bucket/input_b92b9b89-4585-4146-8ee5-22fe99802a8e.json",
-                output_gcs_uri="gs://my-input-bucket/output_b92b9b89-4585-4146-8ee5-22fe99802a8e.json",
+                input_gcs_uri="gs://my-input-bucket/b92b9b89-4585-4146-8ee5-22fe99802a8e_input.json",
+                output_gcs_uri="gs://my-input-bucket/b92b9b89-4585-4146-8ee5-22fe99802a8e_output.json",
             )
 
             request_mock.assert_called_with(
                 "post",
                 f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:asyncQuery",
                 {
                     "_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME},
-                    "inputGcsUri": "gs://my-input-bucket/input_b92b9b89-4585-4146-8ee5-22fe99802a8e.json",
-                    "outputGcsUri": "gs://my-input-bucket/output_b92b9b89-4585-4146-8ee5-22fe99802a8e.json",
+                    "inputGcsUri": "gs://my-input-bucket/b92b9b89-4585-4146-8ee5-22fe99802a8e_input.json",
+                    "outputGcsUri": "gs://my-input-bucket/b92b9b89-4585-4146-8ee5-22fe99802a8e_output.json",
                 },
                 None,
             )
@@ -2980,38 +2980,18 @@ def test_run_query_job_agent_engine_missing_query(self):
         ):
             self.client.agent_engines.run_query_job(
                 name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
-                config={"gcs_bucket": "gs://my-input-bucket/"},
+                config={"output_gcs_uri": "gs://my-input-bucket/"},
             )
 
-    def test_run_query_job_agent_engine_missing_bucket(self):
+    def test_run_query_job_agent_engine_missing_uri(self):
         with pytest.raises(
-            ValueError, match="`gcs_bucket` is required in the config object."
+            ValueError, match="`output_gcs_uri` is required in the config object."
         ):
             self.client.agent_engines.run_query_job(
                 name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
                 config={"query": _TEST_QUERY_PROMPT},
             )
 
-    @mock.patch.object(agent_engines.AgentEngines, "_get")
-    def test_run_query_job_agent_engine_missing_cloud_run_job(self, get_mock):
-        get_mock.return_value = _genai_types.ReasoningEngine(
-            name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
-            spec=_genai_types.ReasoningEngineSpec(
-                deployment_spec=_genai_types.ReasoningEngineSpecDeploymentSpec(env=[])
-            ),
-        )
-        with pytest.raises(
-            ValueError,
-            match="Your ReasoningEngine does not support long running queries, please update your ReasoningEngine and try again.",
-        ):
-            self.client.agent_engines.run_query_job(
-                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
-                config={
-                    "query": _TEST_QUERY_PROMPT,
-                    "gcs_bucket": "gs://my-input-bucket/",
-                },
-            )
-
     @mock.patch("google.cloud.storage.Client")
     @mock.patch.object(agent_engines.AgentEngines, "_get")
     @mock.patch("uuid.uuid4")
@@ -3053,10 +3033,103 @@ def test_run_query_job_agent_engine_bucket_creation_forbidden(
                     name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
                     config={
                         "query": _TEST_QUERY_PROMPT,
-                        "gcs_bucket": "gs://my-input-bucket/",
+                        "output_gcs_uri": "gs://my-input-bucket/",
                     },
                 )
 
+    @mock.patch("google.cloud.storage.Client")
+    @mock.patch.object(agent_engines.AgentEngines, "_get")
+    @mock.patch("uuid.uuid4")
+    def test_run_query_job_agent_engine_file_uri(
+        self, mock_uuid, get_mock, mock_storage_client
+    ):
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(
+                body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}'
+            )
+
+            mock_bucket = mock.Mock()
+            mock_bucket.exists.return_value = True
+            mock_blob = mock.Mock()
+            mock_bucket.blob.return_value = mock_blob
+            mock_storage_client.return_value.bucket.return_value = mock_bucket
+
+            get_mock.return_value = _genai_types.ReasoningEngine(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                spec=_genai_types.ReasoningEngineSpec(
+                    deployment_spec=_genai_types.ReasoningEngineSpecDeploymentSpec(
+                        env=[_genai_types.EnvVar(name="input_gcs_uri", value="")]
+                    )
+                ),
+            )
+
+            result = self.client.agent_engines.run_query_job(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                config={
+                    "query": _TEST_QUERY_PROMPT,
+                    "output_gcs_uri": "gs://my-input-bucket/path/output.json",
+                },
+            )
+
+            mock_blob.upload_from_string.assert_called_once_with(_TEST_QUERY_PROMPT)
+            mock_bucket.blob.assert_called_with("path/output_input.json")
+
+            assert result == _genai_types.RunQueryJobResult(
+                job_name="projects/123/locations/us-central1/reasoningEngines/456/operations/789",
+                input_gcs_uri="gs://my-input-bucket/path/output_input.json",
+                output_gcs_uri="gs://my-input-bucket/path/output.json",
+            )
+
+    @mock.patch("google.cloud.storage.Client")
+    @mock.patch.object(agent_engines.AgentEngines, "_get")
+    @mock.patch("uuid.uuid4")
+    def test_run_query_job_agent_engine_directory_no_slash(
+        self, mock_uuid, get_mock, mock_storage_client
+    ):
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(
+                body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}'
+            )
+
+            mock_bucket = mock.Mock()
+            mock_bucket.exists.return_value = True
+            mock_blob = mock.Mock()
+            mock_bucket.blob.return_value = mock_blob
+            mock_storage_client.return_value.bucket.return_value = mock_bucket
+
+            mock_uuid.return_value.hex = "b92b9b89-4585-4146-8ee5-22fe99802a8e"
+
+            get_mock.return_value = _genai_types.ReasoningEngine(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                spec=_genai_types.ReasoningEngineSpec(
+                    deployment_spec=_genai_types.ReasoningEngineSpecDeploymentSpec(
+                        env=[_genai_types.EnvVar(name="input_gcs_uri", value="")]
+                    )
+                ),
+            )
+
+            result = self.client.agent_engines.run_query_job(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                config={
+                    "query": _TEST_QUERY_PROMPT,
+                    "output_gcs_uri": "gs://my-input-bucket/path",
+                },
+            )
+
+            mock_bucket.blob.assert_called_with(
+                "path/b92b9b89-4585-4146-8ee5-22fe99802a8e_input.json"
+            )
+
+            assert result == _genai_types.RunQueryJobResult(
+                job_name="projects/123/locations/us-central1/reasoningEngines/456/operations/789",
+                input_gcs_uri="gs://my-input-bucket/path/b92b9b89-4585-4146-8ee5-22fe99802a8e_input.json",
+                output_gcs_uri="gs://my-input-bucket/path/b92b9b89-4585-4146-8ee5-22fe99802a8e_output.json",
+            )
+
     def test_query_agent_engine_async(self):
         agent = self.client.agent_engines._register_api_methods(
             agent_engine=_genai_types.AgentEngine(

vertexai/_genai/agent_engines.py

@@ -1077,7 +1077,7 @@ def run_query_job(
                 the default configuration will be used. This can be used to specify
                 the following fields:
                   - query: The query to send to the agent engine.
-                  - gcs_bucket: The GCS bucket path to use for the query.
+                  - output_gcs_uri: The GCS URI to use for the output.
         """
         from google.cloud import storage  # type: ignore[attr-defined]
         from google.api_core import exceptions
@@ -1090,41 +1090,40 @@ def run_query_job(
 
         if not config.query:
             raise ValueError("`query` is required in the config object.")
-        if not config.gcs_bucket:
-            raise ValueError("`gcs_bucket` is required in the config object.")
-
-        api_resource = self._get(name=name)
-
-        is_supported = False
-        if (
-            api_resource.spec
-            and api_resource.spec.deployment_spec
-            and api_resource.spec.deployment_spec.env
-        ):
-            for env in api_resource.spec.deployment_spec.env:
-                if env.name in [
-                    "INPUT_GCS_URI",
-                    "OUTPUT_GCS_URI",
-                    "input_gcs_uri",
-                    "output_gcs_uri",
-                ]:
-                    is_supported = True
-                    break
-
-        if not is_supported:
-            raise ValueError(
-                "Your ReasoningEngine does not support long running queries, "
-                "please update your ReasoningEngine and try again."
-            )
-
-        gcs_bucket = config.gcs_bucket.rstrip("/")
+        if not config.output_gcs_uri:
+            raise ValueError("`output_gcs_uri` is required in the config object.")
+
+        output_gcs_uri = config.output_gcs_uri
+        is_file = False
+        last_part = ""
+        if not output_gcs_uri.endswith("/"):
+            last_part = output_gcs_uri.split("/")[-1]
+            if "." in last_part:
+                is_file = True
+
+        if is_file:
+            path_parts = output_gcs_uri.split("/")
+            file_name = path_parts[-1]
+            base_uri = "/".join(path_parts[:-1])
+            name_parts = file_name.rsplit(".", 1)
+            if len(name_parts) == 2:
+                name_part, ext = name_parts[0], "." + name_parts[1]
+            else:
+                name_part = name_parts[0]
+                ext = ""
+            input_gcs_uri = f"{base_uri}/{name_part}_input{ext}"
+        else:
+            job_uuid = uuid.uuid4().hex
+            gcs_path = output_gcs_uri.rstrip("/")
+            input_gcs_uri = f"{gcs_path}/{job_uuid}_input.json"
+            output_gcs_uri = f"{gcs_path}/{job_uuid}_output.json"
 
         storage_client = storage.Client(
             project=self._api_client.project, credentials=self._api_client._credentials
         )
 
         # Handle creating the bucket if it does not exist
-        bucket_name = gcs_bucket.replace("gs://", "").split("/")[0]
+        bucket_name = config.output_gcs_uri.replace("gs://", "").split("/")[0]
         bucket = storage_client.bucket(bucket_name)
 
         try:
@@ -1144,15 +1143,10 @@ def run_query_job(
                     "The service account may lack 'storage.buckets.create' permission."
                 ) from e
 
-        job_uuid = uuid.uuid4().hex
-        input_blob_name = f"input_{job_uuid}.json"
-        input_gcs_uri = f"{gcs_bucket}/{input_blob_name}"
+        input_blob_name = input_gcs_uri.replace(f"gs://{bucket_name}/", "")
         blob = bucket.blob(input_blob_name)
         blob.upload_from_string(config.query)
 
-        output_blob_name = f"output_{job_uuid}.json"
-        output_gcs_uri = f"{gcs_bucket}/{output_blob_name}"
-
         new_config = types._RunQueryJobAgentEngineConfig(
             input_gcs_uri=input_gcs_uri,
             output_gcs_uri=output_gcs_uri,

vertexai/_genai/types/common.py

@@ -15736,8 +15736,14 @@ class RunQueryJobAgentEngineConfig(_common.BaseModel):
     query: Optional[str] = Field(
         default=None, description="""The query to send to the agent engine."""
     )
-    gcs_bucket: Optional[str] = Field(
-        default=None, description="""The GCS bucket to use for the query."""
+    output_gcs_uri: Optional[str] = Field(
+        default=None,
+        description="""The GCS URI to use for the output.
+      If it is a file, the system use this file to store the response.
+      If it represents a directory, the system automatically generate a file
+      for the response.
+      In both cases, the input query will be stored in the same directory under
+      the same file name prefix as the output file.""",
     )
 
 
@@ -15750,8 +15756,13 @@ class RunQueryJobAgentEngineConfigDict(TypedDict, total=False):
     query: Optional[str]
     """The query to send to the agent engine."""
 
-    gcs_bucket: Optional[str]
-    """The GCS bucket to use for the query."""
+    output_gcs_uri: Optional[str]
+    """The GCS URI to use for the output.
+      If it is a file, the system use this file to store the response.
+      If it represents a directory, the system automatically generate a file
+      for the response.
+      In both cases, the input query will be stored in the same directory under
+      the same file name prefix as the output file."""
 
 
 RunQueryJobAgentEngineConfigOrDict = Union[