fix: Rag - fix upload file environment mismatch

PiperOrigin-RevId: 703549419

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -163,8 +163,10 @@
 TEST_FILE_DISPLAY_NAME = "my-file.txt"
 TEST_FILE_DESCRIPTION = "my file."
 TEST_HEADERS = {"X-Goog-Upload-Protocol": "multipart"}
-TEST_UPLOAD_REQUEST_URI = "https://{}/upload/v1/projects/{}/locations/{}/ragCorpora/{}/ragFiles:upload".format(
-    TEST_API_ENDPOINT, TEST_PROJECT_NUMBER, TEST_REGION, TEST_RAG_CORPUS_ID
+TEST_UPLOAD_REQUEST_URI = (
+    "https://{}/v1/projects/{}/locations/{}/ragCorpora/{}/ragFiles:upload".format(
+        TEST_API_ENDPOINT, TEST_PROJECT_NUMBER, TEST_REGION, TEST_RAG_CORPUS_ID
+    )
 )
 TEST_RAG_FILE_ID = "generate-456"
 TEST_RAG_FILE_RESOURCE_NAME = (

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -308,8 +308,10 @@
 TEST_FILE_DISPLAY_NAME = "my-file.txt"
 TEST_FILE_DESCRIPTION = "my file."
 TEST_HEADERS = {"X-Goog-Upload-Protocol": "multipart"}
-TEST_UPLOAD_REQUEST_URI = "https://{}/upload/v1beta1/projects/{}/locations/{}/ragCorpora/{}/ragFiles:upload".format(
-    TEST_API_ENDPOINT, TEST_PROJECT_NUMBER, TEST_REGION, TEST_RAG_CORPUS_ID
+TEST_UPLOAD_REQUEST_URI = (
+    "https://{}/v1beta1/projects/{}/locations/{}/ragCorpora/{}/ragFiles:upload".format(
+        TEST_API_ENDPOINT, TEST_PROJECT_NUMBER, TEST_REGION, TEST_RAG_CORPUS_ID
+    )
 )
 TEST_RAG_FILE_ID = "generate-456"
 TEST_RAG_FILE_RESOURCE_NAME = (

vertexai/preview/rag/rag_data.py

@@ -409,9 +409,14 @@ def upload_file(
     if display_name is None:
         display_name = "vertex-" + utils.timestamped_unique_name()
     headers = {"X-Goog-Upload-Protocol": "multipart"}
-    upload_request_uri = "https://{}-{}/upload/v1beta1/{}/ragFiles:upload".format(
-        location,
-        aiplatform.constants.base.API_BASE_PATH,
+    if not initializer.global_config.api_endpoint:
+        request_endpoint = "{}-{}".format(
+            location, aiplatform.constants.base.API_BASE_PATH
+        )
+    else:
+        request_endpoint = initializer.global_config.api_endpoint
+    upload_request_uri = "https://{}/v1beta1/{}/ragFiles:upload".format(
+        request_endpoint,
         corpus_name,
     )
     js_rag_file = {"rag_file": {"display_name": display_name}}
@@ -448,7 +453,9 @@ def upload_file(
         raise RuntimeError("Failed in uploading the RagFile due to: ", e) from e
 
     if response.status_code == 404:
-        raise ValueError("RagCorpus '%s' is not found.", corpus_name)
+        raise ValueError(
+            "RagCorpus '%s' is not found: %s", corpus_name, upload_request_uri
+        )
     if response.json().get("error"):
         raise RuntimeError(
             "Failed in indexing the RagFile due to: ", response.json().get("error")

vertexai/rag/rag_data.py

@@ -332,9 +332,14 @@ def upload_file(
     if display_name is None:
         display_name = "vertex-" + utils.timestamped_unique_name()
     headers = {"X-Goog-Upload-Protocol": "multipart"}
-    upload_request_uri = "https://{}-{}/upload/v1/{}/ragFiles:upload".format(
-        location,
-        aiplatform.constants.base.API_BASE_PATH,
+    if not initializer.global_config.api_endpoint:
+        request_endpoint = "{}-{}".format(
+            location, aiplatform.constants.base.API_BASE_PATH
+        )
+    else:
+        request_endpoint = initializer.global_config.api_endpoint
+    upload_request_uri = "https://{}/v1/{}/ragFiles:upload".format(
+        request_endpoint,
         corpus_name,
     )
     js_rag_file = {"rag_file": {"display_name": display_name}}
@@ -372,7 +377,9 @@ def upload_file(
         raise RuntimeError("Failed in uploading the RagFile due to: ", e) from e
 
     if response.status_code == 404:
-        raise ValueError("RagCorpus '%s' is not found.", corpus_name)
+        raise ValueError(
+            "RagCorpus '%s' is not found: %s", corpus_name, upload_request_uri
+        )
     if response.json().get("error"):
         raise RuntimeError(
             "Failed in indexing the RagFile due to: ", response.json().get("error")