feat: ecs batchrag - download encrypted artifacts using auth headers (#1249)

Author: gcuip

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.6.29"
+version = "2.6.30"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/platform/context_grounding/_context_grounding_service.py

@@ -782,9 +782,20 @@ def download_batch_transform_result(
 
         Path(destination_path).parent.mkdir(parents=True, exist_ok=True)
 
+        # SAS uris can be downloaded without authentication
+        # encrypted artifacts require authenticated DownloadBlob endpoint
         with open(destination_path, "wb") as file:
-            with httpx.Client(**get_httpx_client_kwargs()) as client:
-                file_content = client.get(uri_response.uri).content
+            if uri_response.is_encrypted:
+                download_spec = self._batch_transform_download_blob_spec(id=id)
+                download_response = self.request(
+                    download_spec.method,
+                    download_spec.endpoint,
+                    headers=download_spec.headers,
+                )
+                file_content = download_response.content
+            else:
+                with httpx.Client(**get_httpx_client_kwargs()) as client:
+                    file_content = client.get(uri_response.uri).content
             file.write(file_content)
 
     @resource_override(resource_type="index", resource_identifier="index_name")
@@ -828,9 +839,20 @@ async def download_batch_transform_result_async(
         )
         uri_response = BatchTransformReadUriResponse.model_validate(response.json())
 
-        async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
-            download_response = await client.get(uri_response.uri)
+        # SAS uris can be downloaded without authentication
+        # encrypted artifacts require authenticated DownloadBlob endpoint
+        if uri_response.is_encrypted:
+            download_spec = self._batch_transform_download_blob_spec(id=id)
+            download_response = await self.request_async(
+                download_spec.method,
+                download_spec.endpoint,
+                headers=download_spec.headers,
+            )
             file_content = download_response.content
+        else:
+            async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client:
+                download_response = await client.get(uri_response.uri)
+                file_content = download_response.content
 
         Path(destination_path).parent.mkdir(parents=True, exist_ok=True)
 
@@ -1515,6 +1537,15 @@ def _batch_transform_get_read_uri_spec(
             endpoint=Endpoint(f"/ecs_/v2/batchRag/{id}/GetReadUri"),
         )
 
+    def _batch_transform_download_blob_spec(
+        self,
+        id: str,
+    ) -> RequestSpec:
+        return RequestSpec(
+            method="GET",
+            endpoint=Endpoint(f"/ecs_/v2/batchRag/{id}/DownloadBlob"),
+        )
+
     def _resolve_folder_key(self, folder_key, folder_path):
         if folder_key is None and folder_path is not None:
             folder_key = self._folders_service.retrieve_key(folder_path=folder_path)

src/uipath/platform/context_grounding/context_grounding.py

@@ -157,6 +157,7 @@ class BatchTransformReadUriResponse(BaseModel):
         arbitrary_types_allowed=True,
     )
     uri: str
+    is_encrypted: bool = Field(alias="isEncrypted", default=False)
 
 
 class DeepRagCreationResponse(BaseModel):

tests/sdk/services/test_context_grounding_service.py

@@ -1597,6 +1597,7 @@ def test_download_batch_transform_result(
             status_code=200,
             json={
                 "uri": "https://storage.example.com/result.csv",
+                "isEncrypted": False,
             },
         )
 
@@ -1670,6 +1671,7 @@ async def test_download_batch_transform_result_async(
             status_code=200,
             json={
                 "uri": "https://storage.example.com/result.csv",
+                "isEncrypted": False,
             },
         )
 
@@ -1741,6 +1743,7 @@ def test_download_batch_transform_result_creates_nested_directories(
             status_code=200,
             json={
                 "uri": "https://storage.example.com/result.csv",
+                "isEncrypted": False,
             },
         )
 
@@ -1760,6 +1763,68 @@ def test_download_batch_transform_result_creates_nested_directories(
         assert destination.read_bytes() == b"col1,col2\nval1,val2"
         assert destination.parent.exists()
 
+    def test_download_batch_transform_result_encrypted(
+        self,
+        httpx_mock: HTTPXMock,
+        service: ContextGroundingService,
+        base_url: str,
+        org: str,
+        tenant: str,
+        version: str,
+        tmp_path,
+    ) -> None:
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id",
+            status_code=200,
+            json={
+                "id": "test-batch-id",
+                "name": "test-batch-transform",
+                "lastBatchRagStatus": "Successful",
+                "prompt": "Summarize documents",
+                "targetFileGlobPattern": "**",
+                "useWebSearchGrounding": False,
+                "outputColumns": [
+                    {"name": "summary", "description": "Document summary"}
+                ],
+                "createdDate": "2024-01-15T10:30:00Z",
+            },
+        )
+
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/GetReadUri",
+            status_code=200,
+            json={
+                "uri": f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob",
+                "isEncrypted": True,
+            },
+        )
+
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob",
+            status_code=200,
+            content=b"encrypted,data\nval1,val2",
+        )
+
+        destination = tmp_path / "result_encrypted.csv"
+        service.download_batch_transform_result(
+            id="test-batch-id",
+            destination_path=str(destination),
+        )
+
+        assert destination.exists()
+        assert destination.read_bytes() == b"encrypted,data\nval1,val2"
+
+        sent_requests = httpx_mock.get_requests()
+        if sent_requests is None:
+            raise Exception("No request was sent")
+
+        # Verify the DownloadBlob endpoint was called with Authorization header
+        download_request = sent_requests[2]
+        assert download_request.method == "GET"
+        assert "/DownloadBlob" in str(download_request.url)
+        assert "Authorization" in download_request.headers
+        assert download_request.headers["Authorization"].startswith("Bearer ")
+
     def test_create_ephemeral_index(
         self,
         httpx_mock: HTTPXMock,
@@ -1903,6 +1968,7 @@ async def test_download_batch_transform_result_async_creates_nested_directories(
             status_code=200,
             json={
                 "uri": "https://storage.example.com/result.csv",
+                "isEncrypted": False,
             },
         )
 
@@ -1921,3 +1987,66 @@ async def test_download_batch_transform_result_async_creates_nested_directories(
         assert destination.exists()
         assert destination.read_bytes() == b"col1,col2\nval1,val2"
         assert destination.parent.exists()
+
+    @pytest.mark.anyio
+    async def test_download_batch_transform_result_async_encrypted(
+        self,
+        httpx_mock: HTTPXMock,
+        service: ContextGroundingService,
+        base_url: str,
+        org: str,
+        tenant: str,
+        version: str,
+        tmp_path,
+    ) -> None:
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id",
+            status_code=200,
+            json={
+                "id": "test-batch-id",
+                "name": "test-batch-transform",
+                "lastBatchRagStatus": "Successful",
+                "prompt": "Summarize documents",
+                "targetFileGlobPattern": "**",
+                "useWebSearchGrounding": False,
+                "outputColumns": [
+                    {"name": "summary", "description": "Document summary"}
+                ],
+                "createdDate": "2024-01-15T10:30:00Z",
+            },
+        )
+
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/GetReadUri",
+            status_code=200,
+            json={
+                "uri": f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob",
+                "isEncrypted": True,
+            },
+        )
+
+        httpx_mock.add_response(
+            url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob",
+            status_code=200,
+            content=b"encrypted,data\nval1,val2",
+        )
+
+        destination = tmp_path / "result_encrypted.csv"
+        await service.download_batch_transform_result_async(
+            id="test-batch-id",
+            destination_path=str(destination),
+        )
+
+        assert destination.exists()
+        assert destination.read_bytes() == b"encrypted,data\nval1,val2"
+
+        sent_requests = httpx_mock.get_requests()
+        if sent_requests is None:
+            raise Exception("No request was sent")
+
+        # Verify the DownloadBlob endpoint was called with Authorization header
+        download_request = sent_requests[2]
+        assert download_request.method == "GET"
+        assert "/DownloadBlob" in str(download_request.url)
+        assert "Authorization" in download_request.headers
+        assert download_request.headers["Authorization"].startswith("Bearer ")