sarvamai 0.1.22a4__py3-none-any.whl → 0.1.22a7__py3-none-any.whl

sarvamai/doc_digitization_job/job.py

@@ -0,0 +1,496 @@
+"""
+Document Digitization Job helper classes.
+
+This module provides high-level abstractions for managing document digitization jobs,
+including file uploads, status polling, and output downloads.
+"""
+
+import asyncio
+import mimetypes
+import os
+import time
+import typing
+
+import httpx
+from http import HTTPStatus
+
+from ..types.doc_digitization_job_status_response import (
+    DocDigitizationJobStatusResponse,
+)
+
+if typing.TYPE_CHECKING:
+    from .client import AsyncDocDigitizationJobClient, DocDigitizationJobClient
+
+
+class AsyncDocDigitizationJob:
+    """
+    Async helper class for managing a Document Digitization job.
+
+    Provides high-level methods for uploading files, starting the job,
+    polling for completion, and downloading outputs.
+    """
+
+    def __init__(self, job_id: str, client: "AsyncDocDigitizationJobClient"):
+        """
+        Initialize the asynchronous document digitization job.
+
+        Parameters
+        ----------
+        job_id : str
+            The unique identifier for the job.
+        client : AsyncDocDigitizationJobClient
+            The async client instance to use for API calls.
+        """
+        self._job_id = job_id
+        self._client = client
+
+    @property
+    def job_id(self) -> str:
+        """Returns the job ID associated with this job instance."""
+        return self._job_id
+
+    async def upload_file(self, file_path: str, timeout: float = 120.0) -> bool:
+        """
+        Upload the input PDF or ZIP file for the document digitization job.
+
+        Parameters
+        ----------
+        file_path : str
+            Path to the PDF or ZIP file to upload.
+        timeout : float, default=120.0
+            Timeout in seconds for the upload request.
+
+        Returns
+        -------
+        bool
+            True if upload was successful.
+
+        Raises
+        ------
+        ValueError
+            If the file is not a PDF or ZIP file.
+        FileNotFoundError
+            If the file does not exist.
+        httpx.HTTPStatusError
+            If the upload request fails.
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext not in [".pdf", ".zip"]:
+            raise ValueError(f"Only PDF or ZIP files are supported, got: {ext}")
+
+        file_name = os.path.basename(file_path)
+
+        # Get upload URL
+        upload_response = await self._client.get_upload_links(
+            job_id=self._job_id,
+            files=[file_name],
+        )
+
+        if not upload_response.upload_urls:
+            raise ValueError("No upload URL returned from API")
+
+        upload_detail = list(upload_response.upload_urls.values())[0]
+        upload_url = upload_detail.file_url
+
+        # Determine content type
+        content_type, _ = mimetypes.guess_type(file_path)
+        if content_type is None:
+            content_type = "application/octet-stream"
+
+        # Upload the file
+        async with httpx.AsyncClient(timeout=timeout) as http_client:
+            with open(file_path, "rb") as f:
+                file_content = f.read()
+
+            response = await http_client.put(
+                upload_url,
+                content=file_content,
+                headers={
+                    "Content-Type": content_type,
+                    "x-ms-blob-type": "BlockBlob",
+                },
+            )
+            response.raise_for_status()
+
+        return True
+
+    async def start(self) -> DocDigitizationJobStatusResponse:
+        """
+        Start the document digitization job processing.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The job status after starting.
+        """
+        return await self._client.start(job_id=self._job_id)
+
+    async def get_status(self) -> DocDigitizationJobStatusResponse:
+        """
+        Retrieve the current status of the job.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The current job status.
+        """
+        return await self._client.get_status(self._job_id)
+
+    async def wait_until_complete(
+        self, poll_interval: int = 5, timeout: int = 1800
+    ) -> DocDigitizationJobStatusResponse:
+        """
+        Polls job status until it completes, partially completes, or fails.
+
+        Parameters
+        ----------
+        poll_interval : int, default=5
+            Seconds between status checks.
+        timeout : int, default=1800
+            Maximum seconds to wait before raising TimeoutError.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The final job status.
+
+        Raises
+        ------
+        TimeoutError
+            If the job does not complete within the timeout period.
+        """
+        start_time = time.time()
+        terminal_states = {"completed", "partiallycompleted", "failed"}
+
+        while True:
+            status = await self.get_status()
+
+            if status.job_state and status.job_state.lower() in terminal_states:
+                return status
+
+            elapsed = time.time() - start_time
+            if elapsed >= timeout:
+                raise TimeoutError(
+                    f"Job {self._job_id} did not complete within {timeout} seconds. "
+                    f"Current state: {status.job_state}"
+                )
+
+            await asyncio.sleep(poll_interval)
+
+    async def download_output(self, output_path: str, timeout: float = 120.0) -> bool:
+        """
+        Download the output file (HTML or Markdown) to the specified path.
+
+        Parameters
+        ----------
+        output_path : str
+            Path where the output file should be saved.
+        timeout : float, default=120.0
+            Timeout in seconds for the download request.
+
+        Returns
+        -------
+        bool
+            True if download was successful.
+
+        Raises
+        ------
+        ValueError
+            If no download URL is available.
+        httpx.HTTPStatusError
+            If the download request fails.
+        """
+        download_response = await self._client.get_download_links(job_id=self._job_id)
+
+        if not download_response.download_urls:
+            raise ValueError("No download URL returned from API")
+
+        download_detail = list(download_response.download_urls.values())[0]
+        download_url = download_detail.file_url
+
+        async with httpx.AsyncClient(timeout=timeout) as http_client:
+            response = await http_client.get(download_url)
+            response.raise_for_status()
+
+            # Ensure output directory exists
+            os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+
+            with open(output_path, "wb") as f:
+                f.write(response.content)
+
+        return True
+
+    async def get_page_metrics(self) -> typing.Optional[typing.Dict[str, int]]:
+        """
+        Get page processing metrics from the job status.
+
+        Returns
+        -------
+        typing.Optional[typing.Dict[str, int]]
+            Dictionary with keys: total_pages, pages_processed, pages_succeeded, pages_failed.
+            Returns None if job_detail is not available.
+        """
+        status = await self.get_status()
+        if status.job_details and len(status.job_details) > 0:
+            detail = status.job_details[0]
+            return {
+                "total_pages": detail.total_pages or 0,
+                "pages_processed": detail.pages_processed or 0,
+                "pages_succeeded": detail.pages_succeeded or 0,
+                "pages_failed": detail.pages_failed or 0,
+            }
+        return None
+
+    async def get_page_errors(self) -> typing.Optional[typing.List[typing.Any]]:
+        """
+        Get any page-level errors from the job.
+
+        Returns
+        -------
+        typing.Optional[typing.List[typing.Any]]
+            List of page errors, or None if not available.
+        """
+        status = await self.get_status()
+        if status.job_details and len(status.job_details) > 0:
+            return status.job_details[0].page_errors
+        return None
+
+
+class DocDigitizationJob:
+    """
+    Sync helper class for managing a Document Digitization job.
+
+    Provides high-level methods for uploading files, starting the job,
+    polling for completion, and downloading outputs.
+    """
+
+    def __init__(self, job_id: str, client: "DocDigitizationJobClient"):
+        """
+        Initialize the document digitization job.
+
+        Parameters
+        ----------
+        job_id : str
+            The unique identifier for the job.
+        client : DocDigitizationJobClient
+            The client instance to use for API calls.
+        """
+        self._job_id = job_id
+        self._client = client
+
+    @property
+    def job_id(self) -> str:
+        """Returns the job ID associated with this job instance."""
+        return self._job_id
+
+    def upload_file(self, file_path: str, timeout: float = 120.0) -> bool:
+        """
+        Upload the input PDF or ZIP file for the document digitization job.
+
+        Parameters
+        ----------
+        file_path : str
+            Path to the PDF or ZIP file to upload.
+        timeout : float, default=120.0
+            Timeout in seconds for the upload request.
+
+        Returns
+        -------
+        bool
+            True if upload was successful.
+
+        Raises
+        ------
+        ValueError
+            If the file is not a PDF or ZIP file.
+        FileNotFoundError
+            If the file does not exist.
+        httpx.HTTPStatusError
+            If the upload request fails.
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext not in [".pdf", ".zip"]:
+            raise ValueError(f"Only PDF or ZIP files are supported, got: {ext}")
+
+        file_name = os.path.basename(file_path)
+
+        # Get upload URL
+        upload_response = self._client.get_upload_links(
+            job_id=self._job_id,
+            files=[file_name],
+        )
+
+        if not upload_response.upload_urls:
+            raise ValueError("No upload URL returned from API")
+
+        upload_detail = list(upload_response.upload_urls.values())[0]
+        upload_url = upload_detail.file_url
+
+        # Determine content type
+        content_type, _ = mimetypes.guess_type(file_path)
+        if content_type is None:
+            content_type = "application/octet-stream"
+
+        # Upload the file
+        with httpx.Client(timeout=timeout) as http_client:
+            with open(file_path, "rb") as f:
+                file_content = f.read()
+
+            response = http_client.put(
+                upload_url,
+                content=file_content,
+                headers={
+                    "Content-Type": content_type,
+                    "x-ms-blob-type": "BlockBlob",
+                },
+            )
+            response.raise_for_status()
+
+        return True
+
+    def start(self) -> DocDigitizationJobStatusResponse:
+        """
+        Start the document digitization job processing.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The job status after starting.
+        """
+        return self._client.start(job_id=self._job_id)
+
+    def get_status(self) -> DocDigitizationJobStatusResponse:
+        """
+        Retrieve the current status of the job.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The current job status.
+        """
+        return self._client.get_status(self._job_id)
+
+    def wait_until_complete(
+        self, poll_interval: int = 5, timeout: int = 1800
+    ) -> DocDigitizationJobStatusResponse:
+        """
+        Polls job status until it completes, partially completes, or fails.
+
+        Parameters
+        ----------
+        poll_interval : int, default=5
+            Seconds between status checks.
+        timeout : int, default=1800
+            Maximum seconds to wait before raising TimeoutError.
+
+        Returns
+        -------
+        DocDigitizationJobStatusResponse
+            The final job status.
+
+        Raises
+        ------
+        TimeoutError
+            If the job does not complete within the timeout period.
+        """
+        start_time = time.time()
+        terminal_states = {"completed", "partiallycompleted", "failed"}
+
+        while True:
+            status = self.get_status()
+
+            if status.job_state and status.job_state.lower() in terminal_states:
+                return status
+
+            elapsed = time.time() - start_time
+            if elapsed >= timeout:
+                raise TimeoutError(
+                    f"Job {self._job_id} did not complete within {timeout} seconds. "
+                    f"Current state: {status.job_state}"
+                )
+
+            time.sleep(poll_interval)
+
+    def download_output(self, output_path: str, timeout: float = 120.0) -> bool:
+        """
+        Download the output file (HTML or Markdown) to the specified path.
+
+        Parameters
+        ----------
+        output_path : str
+            Path where the output file should be saved.
+        timeout : float, default=120.0
+            Timeout in seconds for the download request.
+
+        Returns
+        -------
+        bool
+            True if download was successful.
+
+        Raises
+        ------
+        ValueError
+            If no download URL is available.
+        httpx.HTTPStatusError
+            If the download request fails.
+        """
+        download_response = self._client.get_download_links(job_id=self._job_id)
+
+        if not download_response.download_urls:
+            raise ValueError("No download URL returned from API")
+
+        download_detail = list(download_response.download_urls.values())[0]
+        download_url = download_detail.file_url
+
+        with httpx.Client(timeout=timeout) as http_client:
+            response = http_client.get(download_url)
+            response.raise_for_status()
+
+            # Ensure output directory exists
+            os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+
+            with open(output_path, "wb") as f:
+                f.write(response.content)
+
+        return True
+
+    def get_page_metrics(self) -> typing.Optional[typing.Dict[str, int]]:
+        """
+        Get page processing metrics from the job status.
+
+        Returns
+        -------
+        typing.Optional[typing.Dict[str, int]]
+            Dictionary with keys: total_pages, pages_processed, pages_succeeded, pages_failed.
+            Returns None if job_detail is not available.
+        """
+        status = self.get_status()
+        if status.job_details and len(status.job_details) > 0:
+            detail = status.job_details[0]
+            return {
+                "total_pages": detail.total_pages or 0,
+                "pages_processed": detail.pages_processed or 0,
+                "pages_succeeded": detail.pages_succeeded or 0,
+                "pages_failed": detail.pages_failed or 0,
+            }
+        return None
+
+    def get_page_errors(self) -> typing.Optional[typing.List[typing.Any]]:
+        """
+        Get any page-level errors from the job.
+
+        Returns
+        -------
+        typing.Optional[typing.List[typing.Any]]
+            List of page errors, or None if not available.
+        """
+        status = self.get_status()
+        if status.job_details and len(status.job_details) > 0:
+            return status.job_details[0].page_errors
+        return None