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: 889620568

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/",
+                    "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,12 +2980,12 @@ 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={"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="`gcs_uri` is required in the config object."
         ):
             self.client.agent_engines.run_query_job(
                 name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
@@ -3008,7 +3008,7 @@ def test_run_query_job_agent_engine_missing_cloud_run_job(self, get_mock):
                 name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
                 config={
                     "query": _TEST_QUERY_PROMPT,
-                    "gcs_bucket": "gs://my-input-bucket/",
+                    "gcs_uri": "gs://my-input-bucket/",
                 },
             )
 
@@ -3053,10 +3053,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/",
+                        "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,
+                    "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,
+                    "gcs_uri": "gs://my-input-bucket/path",
+                },
+            )
+
+            mock_bucket.blob.assert_called_with(
+                "path/input_b92b9b89-4585-4146-8ee5-22fe99802a8e.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(