chunkr-ai 0.0.17__py3-none-any.whl → 0.0.19__py3-none-any.whl

chunkr_ai/api/task_response.py

@@ -0,0 +1,119 @@
+from datetime import datetime
+from typing import TypeVar, Optional, Generic
+from pydantic import BaseModel, PrivateAttr
+import asyncio
+
+from .config import Configuration, OutputConfiguration, OutputResponse, Status
+from .protocol import ChunkrClientProtocol
+from .misc import prepare_upload_data
+from .decorators import anywhere, require_task
+
+T = TypeVar("T", bound="TaskResponse")
+
+class TaskResponse(BaseModel, Generic[T]):
+    configuration: OutputConfiguration
+    created_at: datetime
+    expires_at: Optional[datetime]
+    finished_at: Optional[datetime]
+    message: str
+    output: Optional[OutputResponse]
+    started_at: Optional[datetime]
+    status: Status
+    task_id: str
+    task_url: Optional[str]
+    _client: Optional[ChunkrClientProtocol] = PrivateAttr(default=None)
+
+    def with_client(self, client: ChunkrClientProtocol) -> T:
+        self._client = client
+        return self
+
+    def _check_status(self) -> Optional[T]:
+        """Helper method to check task status and handle completion/failure"""
+        if self.status == "Failed":
+            raise ValueError(self.message)
+        if self.status not in ("Starting", "Processing"):
+            return self
+        return None
+
+    async def _poll_request(self) -> dict:
+        try:
+            if not self._client._client:
+                raise ValueError("Client not found")
+            r = await self._client._client.get(
+                self.task_url, headers=self._client._headers()
+            )
+            r.raise_for_status()
+            return r.json()
+        except (ConnectionError, TimeoutError) as _:
+            print("Connection error while polling the task, retrying...")
+            await asyncio.sleep(0.5)
+        except Exception:
+            raise
+
+    @anywhere()
+    @require_task()
+    async def poll(self) -> T:
+        """Poll the task for completion."""
+        while True:
+            j = await self._poll_request()
+            updated = TaskResponse(**j).with_client(self._client)
+            self.__dict__.update(updated.__dict__)
+            if res := self._check_status():
+                return res
+            await asyncio.sleep(0.5)
+
+    @anywhere()
+    @require_task()
+    async def update(self, config: Configuration) -> T:
+        """Update the task configuration."""
+        f = prepare_upload_data(None, config)
+        r = await self._client._client.patch(
+            self.task_url, files=f, headers=self._client._headers()
+        )
+        r.raise_for_status()
+        updated = TaskResponse(**r.json()).with_client(self._client)
+        self.__dict__.update(updated.__dict__)
+        return await self.poll()
+
+    @anywhere()
+    @require_task()
+    async def delete(self) -> T:
+        """Delete the task."""
+        r = await self._client._client.delete(
+            self.task_url, headers=self._client._headers()
+        )
+        r.raise_for_status()
+        return self
+
+    @anywhere()
+    @require_task()
+    async def cancel(self) -> T:
+        """Cancel the task."""
+        r = await self._client._client.get(
+            f"{self.task_url}/cancel", headers=self._client._headers()
+        )
+        r.raise_for_status()
+        return await self.poll()
+
+    def html(self) -> str:
+        """Get the full HTML of the task"""
+        return self._get_content("html")
+
+    def markdown(self) -> str:
+        """Get the full markdown of the task"""
+        return self._get_content("markdown")
+
+    def content(self) -> str:
+        """Get the full content of the task"""
+        return self._get_content("content")
+
+    def _get_content(self, t: str) -> str:
+        if not self.output:
+            return ""
+        parts = []
+        for c in self.output.chunks:
+            for s in c.segments:
+                v = getattr(s, t)
+                if v:
+                    parts.append(v)
+        return "\n".join(parts)

chunkr_ai/models.py

@@ -4,25 +4,20 @@ from .api.config import (
     ChunkProcessing,
     Configuration,
     CroppingStrategy,
-    ExtractedJson,
     GenerationStrategy,
     GenerationConfig,
-    JsonSchema,
     Model,
     OCRResult,
     OcrStrategy,
     OutputResponse,
-    PipelineType,
-    Property,
     Segment,
     SegmentProcessing,
     SegmentType,
     SegmentationStrategy,
     Status,
+    Pipeline,
 )
-
-from .api.task import TaskResponse
-from .api.task_async import TaskResponseAsync
+from .api.task_response import TaskResponse
 
 __all__ = [
     "BoundingBox",
@@ -30,21 +25,17 @@ __all__ = [
     "ChunkProcessing",
     "Configuration",
     "CroppingStrategy",
-    "ExtractedJson",
     "GenerationConfig",
     "GenerationStrategy",
-    "JsonSchema",
     "Model",
     "OCRResult",
     "OcrStrategy",
     "OutputResponse",
-    "PipelineType",
-    "Property",
     "Segment",
     "SegmentProcessing",
     "SegmentType",
     "SegmentationStrategy",
     "Status",
     "TaskResponse",
-    "TaskResponseAsync",
+    "Pipeline",
 ]

{chunkr_ai-0.0.17.dist-info → chunkr_ai-0.0.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.17
+Version: 0.0.19
 Summary: Python client for Chunkr: open source document intelligence
 Author-email: Ishaan Kapoor <ishaan@lumina.sh>
 Project-URL: Homepage, https://chunkr.ai
@@ -11,7 +11,6 @@ Requires-Dist: pillow>=10.0.0
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: pytest-asyncio>=0.21.0
 Requires-Dist: python-dotenv>=0.19.0
-Requires-Dist: requests>=2.28.0
 Provides-Extra: test
 Requires-Dist: pytest>=7.0.0; extra == "test"
 Requires-Dist: pytest-xdist>=3.0.0; extra == "test"
@@ -34,7 +33,7 @@ pip install chunkr-ai
 
 ## Usage
 
-We provide two clients: `Chunkr` for synchronous operations and `ChunkrAsync` for asynchronous operations.
+The `Chunkr` client works seamlessly in both synchronous and asynchronous contexts.
 
 ### Synchronous Usage
 
@@ -46,62 +45,86 @@ chunkr = Chunkr()
 
 # Upload a file and wait for processing
 task = chunkr.upload("document.pdf")
+print(task.task_id)
 
-# Print the response
-print(task)
+# Create task without waiting
+task = chunkr.create_task("document.pdf")
+result = task.poll()  # Check status when needed
 
-# Get output from task
-output = task.output
-
-# If you want to upload without waiting for processing
-task = chunkr.start_upload("document.pdf")
-# ... do other things ...
-task.poll()  # Check status when needed
+# Clean up when done
+chunkr.close()
 ```
 
 ### Asynchronous Usage
 
 ```python
-from chunkr_ai import ChunkrAsync
+from chunkr_ai import Chunkr
+import asyncio
 
 async def process_document():
     # Initialize client
-    chunkr = ChunkrAsync()
+    chunkr = Chunkr()
+
+    try:
+        # Upload a file and wait for processing
+        task = await chunkr.upload("document.pdf")
+        print(task.task_id)
+
+        # Create task without waiting
+        task = await chunkr.create_task("document.pdf")
+        result = await task.poll()  # Check status when needed
+    finally:
+        await chunkr.close()
 
-    # Upload a file and wait for processing
-    task = await chunkr.upload("document.pdf")
+# Run the async function
+asyncio.run(process_document())
+```
 
-    # Print the response
-    print(task)
+### Concurrent Processing
 
-    # Get output from task
-    output = task.output
+The client supports both async concurrency and multiprocessing:
 
-    # If you want to upload without waiting for processing
-    task = await chunkr.start_upload("document.pdf")
-    # ... do other things ...
-    await task.poll()  # Check status when needed
+```python
+# Async concurrency
+async def process_multiple():
+    chunkr = Chunkr()
+    try:
+        tasks = [
+            chunkr.upload("doc1.pdf"),
+            chunkr.upload("doc2.pdf"),
+            chunkr.upload("doc3.pdf")
+        ]
+        results = await asyncio.gather(*tasks)
+    finally:
+        await chunkr.close()
+
+# Multiprocessing
+from multiprocessing import Pool
+
+def process_file(path):
+    chunkr = Chunkr()
+    try:
+        return chunkr.upload(path)
+    finally:
+        chunkr.close()
+
+with Pool(processes=3) as pool:
+    results = pool.map(process_file, ["doc1.pdf", "doc2.pdf", "doc3.pdf"])
 ```
 
-### Additional Features
+### Input Types
 
-Both clients support various input types:
+The client supports various input types:
 
 ```python
-# Upload from file path
+# File path
 chunkr.upload("document.pdf")
 
-# Upload from opened file
+# Opened file
 with open("document.pdf", "rb") as f:
     chunkr.upload(f)
 
-# Upload from URL
-chunkr.upload("https://example.com/document.pdf")
-
-# Upload from base64 string
-chunkr.upload("data:application/pdf;base64,JVBERi0xLjcKCjEgMCBvYmo...")
-
-# Upload an image
+# PIL Image
 from PIL import Image
 img = Image.open("photo.jpg")
 chunkr.upload(img)
@@ -112,9 +135,13 @@ chunkr.upload(img)
 You can customize the processing behavior by passing a `Configuration` object:
 
 ```python
-from chunkr_ai.models import Configuration, OcrStrategy, SegmentationStrategy, GenerationStrategy
+from chunkr_ai.models import (
+    Configuration, 
+    OcrStrategy, 
+    SegmentationStrategy, 
+    GenerationStrategy
+)
 
-# Basic configuration
 config = Configuration(
     ocr_strategy=OcrStrategy.AUTO,
     segmentation_strategy=SegmentationStrategy.LAYOUT_ANALYSIS,
@@ -122,8 +149,9 @@ config = Configuration(
     expires_in=3600,  # seconds
 )
 
-# Upload with configuration
-task = chunkr.upload("document.pdf", config)
+# Works in both sync and async contexts
+task = chunkr.upload("document.pdf", config)  # sync
+task = await chunkr.upload("document.pdf", config)  # async
 ```
 
 #### Available Configuration Examples
@@ -181,7 +209,7 @@ task = chunkr.upload("document.pdf", config)
   )
   ```
 
-## Environment setup
+## Environment Setup
 
 You can provide your API key and URL in several ways:
 1. Environment variables: `CHUNKR_API_KEY` and `CHUNKR_URL`
@@ -193,3 +221,24 @@ chunkr = Chunkr(
     url="https://api.chunkr.ai"
 )
 ```
+
+## Resource Management
+
+It's recommended to properly close the client when you're done:
+
+```python
+# Sync context
+chunkr = Chunkr()
+try:
+    result = chunkr.upload("document.pdf")
+finally:
+    chunkr.close()
+
+# Async context
+async def process():
+    chunkr = Chunkr()
+    try:
+        result = await chunkr.upload("document.pdf")
+    finally:
+        await chunkr.close()
+```

chunkr_ai-0.0.19.dist-info/RECORD

@@ -0,0 +1,17 @@
+chunkr_ai/__init__.py,sha256=6KpYv2lmD6S5z2kc9pqwuLP5VDHmOuu2qDZArUIhb1s,53
+chunkr_ai/models.py,sha256=sEsnoJaL6wz-4R-cYg2WNl6Wmj4Ad_F8B0QuK9t2sZ8,749
+chunkr_ai/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai/api/api.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai/api/auth.py,sha256=hlv0GiUmlsbFO1wLL9sslqOnsBSoBqkL_6Mk2SDvxgE,413
+chunkr_ai/api/chunkr.py,sha256=V56SP8qs7J2QKRCRM9NGlyA1TtDTdFmGYZWbwbFTK_I,2674
+chunkr_ai/api/chunkr_base.py,sha256=TDqEwCCfgshggi_Mzv76FhPj5z21QP8EVj7siczvfao,9826
+chunkr_ai/api/config.py,sha256=NmPTsDvcjkvNx0gNzDTz-oFG5rQC7jm-H70O_crJCw8,4478
+chunkr_ai/api/decorators.py,sha256=PzaTaPBXUMHoSLz6P0sL5JXANFSJff2vjvESKNiOGQY,2566
+chunkr_ai/api/misc.py,sha256=wUG4SpfEEo7NcVK47gmw42dRy9zT5F9S2DtVC4T4ERs,4877
+chunkr_ai/api/protocol.py,sha256=Nt8aWr4ouVwCvoLqVI5vnXJhT2cvxt0sQC-svUk2G5w,458
+chunkr_ai/api/task_response.py,sha256=I0_XJ6WYYu_TwbaSF95wqRPaOm2PhgMKnarxjAx-BZI,3857
+chunkr_ai-0.0.19.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai-0.0.19.dist-info/METADATA,sha256=azezU0CsGwjHNZet5YMKkY8G6yDVWeclXphs3ESTpSw,5696
+chunkr_ai-0.0.19.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+chunkr_ai-0.0.19.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
+chunkr_ai-0.0.19.dist-info/RECORD,,

chunkr_ai/api/base.py

@@ -1,183 +0,0 @@
-from .config import Configuration
-from .task import TaskResponse
-from .auth import HeadersMixin
-from abc import abstractmethod
-from dotenv import load_dotenv
-import io
-import json
-import os
-from pathlib import Path
-from PIL import Image
-import requests
-from typing import BinaryIO, Tuple, Union
-
-
-class ChunkrBase(HeadersMixin):
-    """Base class with shared functionality for Chunkr API clients."""
-
-    def __init__(self, url: str = None, api_key: str = None):
-        load_dotenv()
-        self.url = url or os.getenv("CHUNKR_URL") or "https://api.chunkr.ai"
-        self._api_key = api_key or os.getenv("CHUNKR_API_KEY")
-        if not self._api_key:
-            raise ValueError(
-                "API key must be provided either directly, in .env file, or as CHUNKR_API_KEY environment variable. You can get an api key at: https://www.chunkr.ai"
-            )
-
-        self.url = self.url.rstrip("/")
-
-    def _prepare_file(
-        self, file: Union[str, Path, BinaryIO, Image.Image]
-    ) -> Tuple[str, BinaryIO]:
-        """Convert various file types into a tuple of (filename, file-like object).
-
-        Args:
-            file: Input file, can be:
-                - String or Path to a file
-                - URL string starting with http:// or https://
-                - Base64 string
-                - Opened binary file (mode='rb')
-                - PIL/Pillow Image object
-
-        Returns:
-            Tuple[str, BinaryIO]: (filename, file-like object) ready for upload
-
-        Raises:
-            FileNotFoundError: If the file path doesn't exist
-            TypeError: If the file type is not supported
-            ValueError: If the URL is invalid or unreachable
-            ValueError: If the MIME type is unsupported
-        """
-        # Handle URLs
-        if isinstance(file, str) and (
-            file.startswith("http://") or file.startswith("https://")
-        ):
-            response = requests.get(file)
-            response.raise_for_status()
-            file_obj = io.BytesIO(response.content)
-            filename = Path(file.split("/")[-1]).name or "downloaded_file"
-            return filename, file_obj
-
-        # Handle base64 strings
-        if isinstance(file, str) and "," in file and ";base64," in file:
-            try:
-                # Split header and data
-                header, base64_data = file.split(",", 1)
-                import base64
-
-                file_bytes = base64.b64decode(base64_data)
-                file_obj = io.BytesIO(file_bytes)
-
-                # Try to determine format from header
-                format = "bin"
-                mime_type = header.split(":")[-1].split(";")[0].lower()
-
-                # Map MIME types to file extensions
-                mime_to_ext = {
-                    "application/pdf": "pdf",
-                    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
-                    "application/msword": "doc",
-                    "application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
-                    "application/vnd.ms-powerpoint": "ppt",
-                    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
-                    "application/vnd.ms-excel": "xls",
-                    "image/jpeg": "jpg",
-                    "image/png": "png",
-                    "image/jpg": "jpg",
-                }
-
-                if mime_type in mime_to_ext:
-                    format = mime_to_ext[mime_type]
-                else:
-                    raise ValueError(f"Unsupported MIME type: {mime_type}")
-
-                return f"file.{format}", file_obj
-            except Exception as e:
-                raise ValueError(f"Invalid base64 string: {str(e)}")
-
-        # Handle file paths
-        if isinstance(file, (str, Path)):
-            path = Path(file).resolve()
-            if not path.exists():
-                raise FileNotFoundError(f"File not found: {file}")
-            return path.name, open(path, "rb")
-
-        # Handle PIL Images
-        if isinstance(file, Image.Image):
-            img_byte_arr = io.BytesIO()
-            format = file.format or "PNG"
-            file.save(img_byte_arr, format=format)
-            img_byte_arr.seek(0)
-            return f"image.{format.lower()}", img_byte_arr
-
-        # Handle file-like objects
-        if hasattr(file, "read") and hasattr(file, "seek"):
-            # Try to get the filename from the file object if possible
-            name = (
-                getattr(file, "name", "document")
-                if hasattr(file, "name")
-                else "document"
-            )
-            return Path(name).name, file
-
-        raise TypeError(f"Unsupported file type: {type(file)}")
-
-    def _prepare_upload_data(
-        self,
-        file: Union[str, Path, BinaryIO, Image.Image],
-        config: Configuration = None,
-    ) -> Tuple[dict, dict]:
-        """Prepare files and data dictionaries for upload.
-
-        Args:
-            file: The file to upload
-            config: Optional configuration settings
-
-        Returns:
-            Tuple[dict, dict]: (files dict, data dict) ready for upload
-        """
-        filename, file_obj = self._prepare_file(file)
-        files = {"file": (filename, file_obj)}
-        data = {}
-
-        if config:
-            config_dict = config.model_dump(mode="json", exclude_none=True)
-            for key, value in config_dict.items():
-                if isinstance(value, dict):
-                    files[key] = (None, json.dumps(value), "application/json")
-                else:
-                    data[key] = value
-
-        return files, data
-
-    @abstractmethod
-    def upload(
-        self,
-        file: Union[str, Path, BinaryIO, Image.Image],
-        config: Configuration = None,
-    ) -> TaskResponse:
-        """Upload a file and wait for processing to complete.
-
-        Must be implemented by subclasses.
-        """
-        pass
-
-    @abstractmethod
-    def start_upload(
-        self,
-        file: Union[str, Path, BinaryIO, Image.Image],
-        config: Configuration = None,
-    ) -> TaskResponse:
-        """Upload a file for processing and immediately return the task response.
-
-        Must be implemented by subclasses.
-        """
-        pass
-
-    @abstractmethod
-    def get_task(self, task_id: str) -> TaskResponse:
-        """Get a task response by its ID.
-
-        Must be implemented by subclasses.
-        """
-        pass

chunkr_ai/api/chunkr_async.py

@@ -1,120 +0,0 @@
-from .chunkr_base import ChunkrBase
-from .config import Configuration
-from .misc import prepare_upload_data
-from .task_async import TaskResponseAsync
-import httpx
-from pathlib import Path
-from PIL import Image
-from typing import Union, BinaryIO
-
-
-class ChunkrAsync(ChunkrBase):
-    """Asynchronous Chunkr API client"""
-
-    def __init__(self, url: str = None, api_key: str = None):
-        super().__init__(url, api_key)
-        self._client = httpx.AsyncClient()
-
-    async def upload(
-        self,
-        file: Union[str, Path, BinaryIO, Image.Image],
-        config: Configuration = None,
-    ) -> TaskResponseAsync:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            task = await self.create_task(file, config)
-            return await task.poll()
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def update(self, task_id: str, config: Configuration) -> TaskResponseAsync:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            task = await self.update_task(task_id, config)
-            return await task.poll()
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def create_task(
-        self,
-        file: Union[str, Path, BinaryIO, Image.Image],
-        config: Configuration = None,
-    ) -> TaskResponseAsync:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            files = prepare_upload_data(file, config)
-            r = await self._client.post(
-                f"{self.url}/api/v1/task", files=files, headers=self._headers()
-            )
-            r.raise_for_status()
-            return TaskResponseAsync(**r.json()).with_client(self)
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def update_task(
-        self, task_id: str, config: Configuration
-    ) -> TaskResponseAsync:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            files = prepare_upload_data(None, config)
-            r = await self._client.patch(
-                f"{self.url}/api/v1/task/{task_id}",
-                files=files,
-                headers=self._headers(),
-            )
-
-            r.raise_for_status()
-            return TaskResponseAsync(**r.json()).with_client(self)
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def get_task(self, task_id: str) -> TaskResponseAsync:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            r = await self._client.get(
-                f"{self.url}/api/v1/task/{task_id}", headers=self._headers()
-            )
-            r.raise_for_status()
-            return TaskResponseAsync(**r.json()).with_client(self)
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def delete_task(self, task_id: str) -> None:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            r = await self._client.delete(
-                f"{self.url}/api/v1/task/{task_id}", headers=self._headers()
-            )
-            r.raise_for_status()
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def cancel_task(self, task_id: str) -> None:
-        if not self._client or self._client.is_closed:
-            self._client = httpx.AsyncClient()
-        try:
-            r = await self._client.get(
-                f"{self.url}/api/v1/task/{task_id}/cancel", headers=self._headers()
-            )
-            r.raise_for_status()
-        except Exception as e:
-            await self._client.aclose()
-            raise e
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self._client.aclose()