From 6e6cde46014c386fe0c0398b1bdcb22b37d9515e Mon Sep 17 00:00:00 2001 From: Gabriel Martin <91462031+gcuip@users.noreply.github.com> Date: Mon, 2 Feb 2026 16:48:26 +0100 Subject: [PATCH 1/2] feat(ecs): batchrag - download encrypted artifacts using auth headers --- .../_context_grounding_service.py | 23 +++- .../context_grounding/context_grounding.py | 1 + .../test_context_grounding_service.py | 127 ++++++++++++++++++ 3 files changed, 147 insertions(+), 4 deletions(-) diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 31cadcc0a..ebca1b9f1 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -783,8 +783,15 @@ def download_batch_transform_result( Path(destination_path).parent.mkdir(parents=True, exist_ok=True) 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: + # Use authenticated client for encrypted artifacts + file_content = self._client.get( + uri_response.uri, headers=self.auth_headers + ).content + else: + # Use unauthenticated client for non-encrypted artifacts + 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 +835,17 @@ 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) + if uri_response.is_encrypted: + # Use authenticated client for encrypted artifacts + download_response = await self._client_async.get( + uri_response.uri, headers=self.auth_headers + ) file_content = download_response.content + else: + # Use unauthenticated client for non-encrypted artifacts + 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) diff --git a/src/uipath/platform/context_grounding/context_grounding.py b/src/uipath/platform/context_grounding/context_grounding.py index 26de4261a..851e2a81a 100644 --- a/src/uipath/platform/context_grounding/context_grounding.py +++ b/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): diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index a746d9e39..2e0786951 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/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,67 @@ 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}/storage/encrypted/result.csv", + "isEncrypted": True, + }, + ) + + httpx_mock.add_response( + url=f"{base_url}{org}{tenant}/storage/encrypted/result.csv", + 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 download request includes Authorization header + download_request = sent_requests[2] + assert download_request.method == "GET" + assert "Authorization" in download_request.headers + assert download_request.headers["Authorization"].startswith("Bearer ") + def test_create_ephemeral_index( self, httpx_mock: HTTPXMock, @@ -1903,6 +1967,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 +1986,65 @@ 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}/storage/encrypted/result.csv", + "isEncrypted": True, + }, + ) + + httpx_mock.add_response( + url=f"{base_url}{org}{tenant}/storage/encrypted/result.csv", + 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 download request includes Authorization header + download_request = sent_requests[2] + assert download_request.method == "GET" + assert "Authorization" in download_request.headers + assert download_request.headers["Authorization"].startswith("Bearer ") From fd7aca937bd2254127997bc8364641ed341c49fa Mon Sep 17 00:00:00 2001 From: Gabriel Martin <91462031+gcuip@users.noreply.github.com> Date: Tue, 3 Feb 2026 14:20:54 +0100 Subject: [PATCH 2/2] chore: version bump --- pyproject.toml | 2 +- .../_context_grounding_service.py | 34 ++++++++++++++----- .../test_context_grounding_service.py | 14 ++++---- uv.lock | 2 +- 4 files changed, 35 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bf024661e..c4da9a00e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath" -version = "2.6.26" +version = "2.6.31" 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" diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index ebca1b9f1..f2ef712e6 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/src/uipath/platform/context_grounding/_context_grounding_service.py @@ -782,14 +782,18 @@ 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: if uri_response.is_encrypted: - # Use authenticated client for encrypted artifacts - file_content = self._client.get( - uri_response.uri, headers=self.auth_headers - ).content + 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: - # Use unauthenticated client for non-encrypted artifacts with httpx.Client(**get_httpx_client_kwargs()) as client: file_content = client.get(uri_response.uri).content file.write(file_content) @@ -835,14 +839,17 @@ async def download_batch_transform_result_async( ) uri_response = BatchTransformReadUriResponse.model_validate(response.json()) + # SAS uris can be downloaded without authentication + # encrypted artifacts require authenticated DownloadBlob endpoint if uri_response.is_encrypted: - # Use authenticated client for encrypted artifacts - download_response = await self._client_async.get( - uri_response.uri, headers=self.auth_headers + 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: - # Use unauthenticated client for non-encrypted artifacts async with httpx.AsyncClient(**get_httpx_client_kwargs()) as client: download_response = await client.get(uri_response.uri) file_content = download_response.content @@ -1530,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) diff --git a/tests/sdk/services/test_context_grounding_service.py b/tests/sdk/services/test_context_grounding_service.py index 2e0786951..6a1e95e61 100644 --- a/tests/sdk/services/test_context_grounding_service.py +++ b/tests/sdk/services/test_context_grounding_service.py @@ -1794,13 +1794,13 @@ def test_download_batch_transform_result_encrypted( url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/GetReadUri", status_code=200, json={ - "uri": f"{base_url}{org}{tenant}/storage/encrypted/result.csv", + "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}/storage/encrypted/result.csv", + url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob", status_code=200, content=b"encrypted,data\nval1,val2", ) @@ -1818,9 +1818,10 @@ def test_download_batch_transform_result_encrypted( if sent_requests is None: raise Exception("No request was sent") - # Verify the download request includes Authorization header + # 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 ") @@ -2019,13 +2020,13 @@ async def test_download_batch_transform_result_async_encrypted( url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/GetReadUri", status_code=200, json={ - "uri": f"{base_url}{org}{tenant}/storage/encrypted/result.csv", + "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}/storage/encrypted/result.csv", + url=f"{base_url}{org}{tenant}/ecs_/v2/batchRag/test-batch-id/DownloadBlob", status_code=200, content=b"encrypted,data\nval1,val2", ) @@ -2043,8 +2044,9 @@ async def test_download_batch_transform_result_async_encrypted( if sent_requests is None: raise Exception("No request was sent") - # Verify the download request includes Authorization header + # 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 ") diff --git a/uv.lock b/uv.lock index fc6d8b446..00ee127ee 100644 --- a/uv.lock +++ b/uv.lock @@ -2491,7 +2491,7 @@ wheels = [ [[package]] name = "uipath" -version = "2.6.26" +version = "2.6.31" source = { editable = "." } dependencies = [ { name = "applicationinsights" },