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.37"
+version = "0.2.38"
 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

@@ -16,6 +16,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,
@@ -24,6 +25,7 @@
     ExtractionResponse,
     ExtractionResponseIXP,
     FileContent,
+    ModelExtractionResponse,
     ProjectType,
     StartExtractionResponse,
     StartExtractionValidationResponse,
@@ -37,6 +39,9 @@
 # Opts in to DU Framework API v2. Otherwise the service calls API v1.1.
 FRAMEWORK_V2_FEATURE_FLAG = "DocumentsFrameworkV2"
 
+# Models deployed to a folder are served by API v2 only, whatever the flag says.
+_MODEL_API_VERSION = "2.0"
+
 # Response values returned as-is: Document Processing contract payloads are
 # PascalCase in every API version, and the digitization DOM is large and unused.
 _UNNORMALIZED_KEYS = frozenset(
@@ -55,6 +60,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 _operation_result(
     response: Dict[str, Any], operation_id: str, operation_name: str
 ) -> Dict[str, Any]:
@@ -1448,6 +1458,158 @@ 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": _MODEL_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": _MODEL_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": _MODEL_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": _MODEL_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": _MODEL_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": _MODEL_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")