diff --git a/pyproject.toml b/pyproject.toml index 3cf146d66..729e4b731 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath/platform/context_grounding/_context_grounding_service.py b/src/uipath/platform/context_grounding/_context_grounding_service.py index 31cadcc0a..f2ef712e6 100644 --- a/src/uipath/platform/context_grounding/_context_grounding_service.py +++ b/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) 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..6a1e95e61 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,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 ") diff --git a/uv.lock b/uv.lock index 5ae9ebc86..5cb6decfb 100644 --- a/uv.lock +++ b/uv.lock @@ -2491,7 +2491,7 @@ wheels = [ [[package]] name = "uipath" -version = "2.6.29" +version = "2.6.30" source = { editable = "." } dependencies = [ { name = "applicationinsights" },