feat(documents): extract with IXP models deployed to a folder

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.2.35"
+version = "0.2.36"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath-platform/src/uipath/platform/documents/__init__.py

@@ -16,6 +16,7 @@
     FieldType,
     FieldValueProjection,
     FileContent,
+    ModelExtractionResponse,
     ProjectType,
     Reference,
     StartExtractionResponse,
@@ -47,4 +48,5 @@
     "StartExtractionResponse",
     "StartOperationResponse",
     "StartExtractionValidationResponse",
+    "ModelExtractionResponse",
 ]

packages/uipath-platform/src/uipath/platform/documents/_documents_service.py

@@ -15,6 +15,7 @@
 from ..common._execution_context import UiPathExecutionContext
 from ..common._folder_context import FolderContext
 from ..common._models import Endpoint
+from ..constants import HEADER_FOLDER_KEY
 from ..errors import OperationFailedException, OperationNotCompleteException
 from .documents import (
     ActionPriority,
@@ -23,6 +24,7 @@
     ExtractionResponse,
     ExtractionResponseIXP,
     FileContent,
+    ModelExtractionResponse,
     ProjectType,
     StartExtractionResponse,
     StartExtractionValidationResponse,
@@ -53,6 +55,11 @@ def _framework_url(*segments: Any) -> Endpoint:
     return Endpoint(f"/du_/api/framework/projects/{path}")
 
 
+def _model_url(model_name: str, *segments: Any) -> Endpoint:
+    path = "/".join(quote(str(segment), safe="") for segment in (model_name, *segments))
+    return Endpoint(f"/du_/api/framework/models/{path}")
+
+
 def _extractor_url(
     project_id: str,
     extractor_id: Optional[str],
@@ -1415,6 +1422,152 @@ async def retrieve_ixp_extraction_result_async(
 
         return ExtractionResponseIXP.model_validate(extraction_response)
 
+    def _model_headers(self, folder_key: Optional[str]) -> Dict[str, str]:
+        # Models deployed to a folder are addressed by folder key only.
+        key = folder_key or self._folder_key
+        if not key:
+            raise ValueError("`folder_key` must be provided outside a job")
+        return {**self._get_common_headers(), HEADER_FOLDER_KEY: key}
+
+    def _wait_for_model_operation(
+        self, url: Endpoint, headers: Dict[str, str], key: str
+    ) -> Dict[str, Any]:
+        def result_getter() -> Tuple[str, Optional[str], Optional[Dict[str, Any]]]:
+            result = self._request_json(
+                "GET", url=url, params={"api-version": API_VERSION}, headers=headers
+            )
+            return result["status"], result.get("error"), result.get(key)
+
+        return self._wait_for_operation(
+            result_getter=result_getter,
+            wait_statuses=["NotStarted", "Running"],
+            success_status="Succeeded",
+        )
+
+    async def _wait_for_model_operation_async(
+        self, url: Endpoint, headers: Dict[str, str], key: str
+    ) -> Dict[str, Any]:
+        async def result_getter() -> Tuple[
+            str, Optional[str], Optional[Dict[str, Any]]
+        ]:
+            result = await self._request_json_async(
+                "GET", url=url, params={"api-version": API_VERSION}, headers=headers
+            )
+            return result["status"], result.get("error"), result.get(key)
+
+        return await self._wait_for_operation_async(
+            result_getter=result_getter,
+            wait_statuses=["NotStarted", "Running"],
+            success_status="Succeeded",
+        )
+
+    @traced(name="documents_extract_with_model", run_type="uipath")
+    def extract_with_model(
+        self,
+        model_name: str,
+        file: Optional[FileContent] = None,
+        file_path: Optional[str] = None,
+        folder_key: Optional[str] = None,
+    ) -> ModelExtractionResponse:
+        """Extract data from a document with a model deployed to a folder.
+
+        Args:
+            model_name (str): Name of the model deployed to the folder (e.g., "invoices-ixp").
+            file (FileContent, optional): The document file to be processed.
+            file_path (str, optional): Path to the document file to be processed.
+            folder_key (str, optional): Key of the folder the model is deployed to. Defaults to the current folder.
+
+        Note:
+            Either `file` or `file_path` must be provided, but not both.
+
+        Returns:
+            ModelExtractionResponse: The extraction response containing the extracted data.
+
+        Examples:
+            ```python
+            extraction_response = uipath.documents.extract_with_model(
+                model_name="invoices-ixp",
+                file_path="path/to/document.pdf",
+            )
+            ```
+        """
+        _exactly_one_must_be_provided(file=file, file_path=file_path)
+        headers = self._model_headers(folder_key)
+
+        with open(Path(file_path), "rb") if file_path else nullcontext(file) as handle:
+            document_id = self._request_json(
+                "POST",
+                url=_model_url(model_name, "digitization", "start"),
+                params={"api-version": API_VERSION},
+                headers=headers,
+                files={"File": handle},
+            )["documentId"]
+        self._wait_for_model_operation(
+            _model_url(model_name, "digitization", "result", document_id),
+            headers,
+            "result",
+        )
+        operation_id = self._request_json(
+            "POST",
+            url=_model_url(model_name, "extraction", "start"),
+            params={"api-version": API_VERSION},
+            headers=headers,
+            json={"documentId": document_id},
+        )["operationId"]
+        result = self._wait_for_model_operation(
+            _model_url(model_name, "extraction", "result", operation_id),
+            headers,
+            "result",
+        )
+        return ModelExtractionResponse.model_validate(
+            {**result, "modelName": model_name}
+        )
+
+    @traced(name="documents_extract_with_model_async", run_type="uipath")
+    async def extract_with_model_async(
+        self,
+        model_name: str,
+        file: Optional[FileContent] = None,
+        file_path: Optional[str] = None,
+        folder_key: Optional[str] = None,
+    ) -> ModelExtractionResponse:
+        """Asynchronous version of the [`extract_with_model`][uipath.platform.documents._documents_service.DocumentsService.extract_with_model] method."""
+        _exactly_one_must_be_provided(file=file, file_path=file_path)
+        headers = self._model_headers(folder_key)
+
+        with open(Path(file_path), "rb") if file_path else nullcontext(file) as handle:
+            document_id = (
+                await self._request_json_async(
+                    "POST",
+                    url=_model_url(model_name, "digitization", "start"),
+                    params={"api-version": API_VERSION},
+                    headers=headers,
+                    files={"File": handle},
+                )
+            )["documentId"]
+        await self._wait_for_model_operation_async(
+            _model_url(model_name, "digitization", "result", document_id),
+            headers,
+            "result",
+        )
+        operation_id = (
+            await self._request_json_async(
+                "POST",
+                url=_model_url(model_name, "extraction", "start"),
+                params={"api-version": API_VERSION},
+                headers=headers,
+                json={"documentId": document_id},
+            )
+        )["operationId"]
+        result = await self._wait_for_model_operation_async(
+            _model_url(model_name, "extraction", "result", operation_id),
+            headers,
+            "result",
+        )
+        return ModelExtractionResponse.model_validate(
+            {**result, "modelName": model_name}
+        )
+
     @traced(name="documents_extract", run_type="uipath")
     def extract(
         self,

packages/uipath-platform/src/uipath/platform/documents/documents.py

@@ -303,3 +303,21 @@ class StartExtractionResponse(StartOperationResponse):
 
 class StartExtractionValidationResponse(StartOperationResponse):
     """A model representing the response from starting an extraction validation operation."""
+
+
+class ModelExtractionResponse(BaseModel):
+    """A model representing the result of an extraction with a model deployed to a folder.
+
+    Attributes:
+        extraction_result (ExtractionResult): The result of the extraction process.
+        model_name (str): The name of the deployed model.
+    """
+
+    model_config = ConfigDict(
+        serialize_by_alias=True,
+        validate_by_alias=True,
+        validate_by_name=True,
+    )
+
+    extraction_result: ExtractionResult = Field(alias="extractionResult")
+    model_name: str = Field(alias="modelName")

packages/uipath-platform/tests/services/test_documents_service.py

@@ -3916,3 +3916,159 @@ async def test_start_ixp_extraction_encodes_tag_in_url(
         # ASSERT
         assert response.operation_id == operation_id
         assert response.tag == "my tag/../v2"
+
+
+class TestDocumentsServiceFolderModels:
+    @pytest.mark.parametrize("mode", ["sync", "async"])
+    @pytest.mark.asyncio
+    async def test_extract_with_model(
+        self,
+        httpx_mock: HTTPXMock,
+        service: DocumentsService,
+        base_url: str,
+        org: str,
+        tenant: str,
+        ixp_extraction_response_v2: dict,  # type: ignore
+        mode: str,
+    ):
+        # ARRANGE
+        document_id = str(uuid4())
+        operation_id = str(uuid4())
+        base = f"{base_url}{org}{tenant}/du_/api/framework/models/invoices-ixp"
+        folder_key = str(uuid4())
+        headers = {
+            "X-UiPath-Internal-ConsumptionSourceType": "CodedAgents",
+            "x-uipath-folderkey": folder_key,
+        }
+        extraction_result = ixp_extraction_response_v2["extractionResult"]
+
+        httpx_mock.add_response(
+            url=f"{base}/digitization/start?api-version=2.0",
+            match_headers=headers,
+            match_files={"File": b"test content"},
+            json={"DocumentId": document_id},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/digitization/result/{document_id}?api-version=2.0",
+            match_headers=headers,
+            json={"Status": "Succeeded", "Result": {}},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/extraction/start?api-version=2.0",
+            match_headers=headers,
+            match_json={"documentId": document_id},
+            json={"OperationId": operation_id},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/extraction/result/{operation_id}?api-version=2.0",
+            match_headers=headers,
+            json={"Status": "Running"},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/extraction/result/{operation_id}?api-version=2.0",
+            match_headers=headers,
+            json={
+                "Status": "Succeeded",
+                "Result": {"ExtractionResult": extraction_result},
+            },
+        )
+
+        # ACT
+        if mode == "async":
+            response = await service.extract_with_model_async(
+                model_name="invoices-ixp", file=b"test content", folder_key=folder_key
+            )
+        else:
+            response = service.extract_with_model(
+                model_name="invoices-ixp", file=b"test content", folder_key=folder_key
+            )
+
+        # ASSERT
+        assert response.model_name == "invoices-ixp"
+        assert response.extraction_result.document_id == extraction_result["DocumentId"]
+        assert (
+            response.extraction_result.results_document
+            == extraction_result["ResultsDocument"]
+        )
+
+    @pytest.mark.parametrize("mode", ["sync", "async"])
+    @pytest.mark.asyncio
+    async def test_extract_with_model_defaults_to_current_folder(
+        self,
+        httpx_mock: HTTPXMock,
+        config: UiPathApiConfig,
+        execution_context: UiPathExecutionContext,
+        monkeypatch: pytest.MonkeyPatch,
+        base_url: str,
+        org: str,
+        tenant: str,
+        ixp_extraction_response_v2: dict,  # type: ignore
+        mode: str,
+    ):
+        # ARRANGE
+        folder_key = str(uuid4())
+        monkeypatch.setenv("UIPATH_FOLDER_KEY", folder_key)
+        service = DocumentsService(
+            config=config, execution_context=execution_context, polling_interval=0
+        )
+        document_id = str(uuid4())
+        operation_id = str(uuid4())
+        base = f"{base_url}{org}{tenant}/du_/api/framework/models/my%20model"
+        headers = {"x-uipath-folderkey": folder_key}
+
+        httpx_mock.add_response(
+            url=f"{base}/digitization/start?api-version=2.0",
+            match_headers=headers,
+            json={"DocumentId": document_id},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/digitization/result/{document_id}?api-version=2.0",
+            match_headers=headers,
+            json={"Status": "Succeeded"},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/extraction/start?api-version=2.0",
+            match_headers=headers,
+            json={"OperationId": operation_id},
+        )
+        httpx_mock.add_response(
+            url=f"{base}/extraction/result/{operation_id}?api-version=2.0",
+            match_headers=headers,
+            json={
+                "Status": "Succeeded",
+                "Result": {
+                    "ExtractionResult": ixp_extraction_response_v2["extractionResult"]
+                },
+            },
+        )
+
+        # ACT
+        if mode == "async":
+            response = await service.extract_with_model_async(
+                model_name="my model", file=b"test content"
+            )
+        else:
+            response = service.extract_with_model(
+                model_name="my model", file=b"test content"
+            )
+
+        # ASSERT
+        assert response.model_name == "my model"
+
+    @pytest.mark.parametrize("mode", ["sync", "async"])
+    @pytest.mark.asyncio
+    async def test_extract_with_model_without_folder(
+        self,
+        service: DocumentsService,
+        mode: str,
+    ):
+        # ACT & ASSERT
+        with pytest.raises(ValueError, match="`folder_key` must be provided"):
+            if mode == "async":
+                await service.extract_with_model_async(
+                    model_name="invoices-ixp", file=b"test content"
+                )
+            else:
+                service.extract_with_model(
+                    model_name="invoices-ixp", file=b"test content"
+                )