chunkr-ai 0.0.23__py3-none-any.whl → 0.0.25__py3-none-any.whl

chunkr_ai/api/chunkr.py

@@ -2,8 +2,8 @@ from pathlib import Path
 from PIL import Image
 from typing import Union, BinaryIO
 
-from .config import Configuration
-from .decorators import anywhere, ensure_client
+from .configuration import Configuration
+from .decorators import anywhere, ensure_client, retry_on_429
 from .misc import prepare_upload_data
 from .task_response import TaskResponse
 from .chunkr_base import ChunkrBase
@@ -29,6 +29,7 @@ class Chunkr(ChunkrBase):
 
     @anywhere()
     @ensure_client()
+    @retry_on_429()
     async def create_task(
         self,
         file: Union[str, Path, BinaryIO, Image.Image],
@@ -39,10 +40,11 @@ class Chunkr(ChunkrBase):
             f"{self.url}/api/v1/task", files=files, headers=self._headers()
         )
         r.raise_for_status()
-        return TaskResponse(**r.json()).with_client(self)
+        return TaskResponse(**r.json()).with_client(self, True, False)
 
     @anywhere()
     @ensure_client()
+    @retry_on_429()
     async def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
         files = await prepare_upload_data(None, config, self._client)
         r = await self._client.patch(
@@ -51,16 +53,22 @@ class Chunkr(ChunkrBase):
             headers=self._headers(),
         )
         r.raise_for_status()
-        return TaskResponse(**r.json()).with_client(self)
+        return TaskResponse(**r.json()).with_client(self, True, False)
 
     @anywhere()
     @ensure_client()
-    async def get_task(self, task_id: str) -> TaskResponse:
+    async def get_task(self, task_id: str, include_chunks: bool = True, base64_urls: bool = False) -> TaskResponse:
+        params = {
+            "base64_urls": str(base64_urls).lower(),
+            "include_chunks": str(include_chunks).lower()
+        }
         r = await self._client.get(
-            f"{self.url}/api/v1/task/{task_id}", headers=self._headers()
+            f"{self.url}/api/v1/task/{task_id}",
+            params=params,
+            headers=self._headers()
         )
         r.raise_for_status()
-        return TaskResponse(**r.json()).with_client(self)
+        return TaskResponse(**r.json()).with_client(self, include_chunks, base64_urls)
 
     @anywhere()
     @ensure_client()

chunkr_ai/api/chunkr_base.py

@@ -1,4 +1,4 @@
-from .config import Configuration
+from .configuration import Configuration
 from .task_response import TaskResponse
 from .auth import HeadersMixin
 from abc import abstractmethod
@@ -11,12 +11,20 @@ from typing import BinaryIO, Union
 
 
 class ChunkrBase(HeadersMixin):
-    """Base class with shared functionality for Chunkr API clients."""
-
-    def __init__(self, url: str = None, api_key: str = None):
+    """Base class with shared functionality for Chunkr API clients.
+    
+    Args:
+        url: The base URL of the Chunkr API. Defaults to the value of the CHUNKR_URL environment variable, or "https://api.chunkr.ai" if not set.
+        api_key: The API key to use for authentication. Defaults to the value of the CHUNKR_API_KEY environment variable, or None if not set.
+        raise_on_failure: Whether to raise an exception if the task fails. Defaults to False.
+    """
+
+    def __init__(self, url: str = None, api_key: str = None, raise_on_failure: bool = False):
         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")
+        self.raise_on_failure = raise_on_failure
+        
         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"
@@ -131,11 +139,13 @@ class ChunkrBase(HeadersMixin):
         pass
 
     @abstractmethod
-    def get_task(self, task_id: str) -> TaskResponse:
+    def get_task(self, task_id: str, include_chunks: bool = True, base64_urls: bool = False) -> TaskResponse:
         """Get a task response by its ID.
 
         Args:
             task_id: The ID of the task to get
+            include_chunks: Whether to include chunks in the output response. Defaults to True.
+            base64_urls: Whether to return base64 encoded URLs. If false, the URLs will be returned as presigned URLs. Defaults to False.
 
         Returns:
             TaskResponse: The task response

chunkr_ai/api/{config.py → configuration.py}

@@ -1,6 +1,6 @@
-from pydantic import BaseModel, Field, model_validator, ConfigDict
+from pydantic import BaseModel, Field, ConfigDict
 from enum import Enum
-from typing import Optional, List, Dict, Union, Type
+from typing import Any, List, Optional
 
 class GenerationStrategy(str, Enum):
     LLM = "LLM"
@@ -37,16 +37,6 @@ class SegmentProcessing(BaseModel):
 class ChunkProcessing(BaseModel):
     target_length: Optional[int] = None
 
-class Property(BaseModel):
-    name: str
-    prop_type: str
-    description: Optional[str] = None
-    default: Optional[str] = None
-
-class JsonSchema(BaseModel):
-    title: str
-    properties: List[Property]
-
 class OcrStrategy(str, Enum):
     ALL = "All"
     AUTO = "Auto"
@@ -98,9 +88,6 @@ class Chunk(BaseModel):
     chunk_length: int
     segments: List[Segment]
 
-class ExtractedJson(BaseModel):
-    data: Dict
-
 class OutputResponse(BaseModel):
     chunks: List[Chunk]
     file_name: Optional[str]
@@ -118,7 +105,6 @@ class Configuration(BaseModel):
     chunk_processing: Optional[ChunkProcessing] = None
     expires_in: Optional[int] = None
     high_resolution: Optional[bool] = None
-    model: Optional[Model] = None
     ocr_strategy: Optional[OcrStrategy] = None
     segment_processing: Optional[SegmentProcessing] = None
     segmentation_strategy: Optional[SegmentationStrategy] = None
@@ -126,16 +112,10 @@ class Configuration(BaseModel):
     
 class OutputConfiguration(Configuration):
     input_file_url: Optional[str] = None
-    json_schema: Optional[Union[JsonSchema, Type[BaseModel], BaseModel]] = None
-    
-    @model_validator(mode="before")
-    def map_deprecated_fields(cls, values: Dict) -> Dict:
-        if isinstance(values, dict) and "target_chunk_length" in values:
-            target_length = values.pop("target_chunk_length")
-            if target_length is not None:
-                values["chunk_processing"] = values.get("chunk_processing", {}) or {}
-                values["chunk_processing"]["target_length"] = target_length
-        return values
+    # Deprecated
+    json_schema: Optional[Any] = None
+    model: Optional[Model] = None
+    target_chunk_length: Optional[int] = None
     
 class Status(str, Enum):
     STARTING = "Starting"

chunkr_ai/api/decorators.py

@@ -59,4 +59,34 @@ def require_task() -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitabl
                 self._client._client = httpx.AsyncClient()
             return await async_func(self, *args, **kwargs)
         return wrapper
+    return decorator
+
+def retry_on_429(max_retries: int = 10, initial_delay: float = 0.5) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+    """Decorator that retries the request when encountering 429 Too Many Requests errors.
+    
+    Args:
+        max_retries: Maximum number of retry attempts (default: 3)
+        initial_delay: Initial delay in seconds, will be exponentially increased (default: 1.0)
+    """
+    def decorator(async_func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+        @functools.wraps(async_func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            retries = 0
+            while True:
+                try:
+                    return await async_func(*args, **kwargs)
+                except httpx.HTTPStatusError as e:
+                    if e.response.status_code != 429 or retries >= max_retries:
+                        raise
+                    retries += 1
+                    delay = initial_delay
+                    # Use Retry-After header if available
+                    retry_after = e.response.headers.get('Retry-After')
+                    if retry_after:
+                        try:
+                            delay = float(retry_after)
+                        except (ValueError, TypeError):
+                            pass
+                    await asyncio.sleep(delay)
+        return wrapper
     return decorator

chunkr_ai/api/misc.py

@@ -1,4 +1,4 @@
-from .config import Configuration
+from .configuration import Configuration
 import io
 import json
 from pathlib import Path

chunkr_ai/api/protocol.py

@@ -5,15 +5,10 @@ from httpx import AsyncClient
 @runtime_checkable
 class ChunkrClientProtocol(Protocol):
     """Protocol defining the interface for Chunkr clients"""
-
-    url: str
-    _api_key: str
+    
+    raise_on_failure: bool = True
     _client: Optional[AsyncClient] = None
 
-    def get_api_key(self) -> str:
-        """Get the API key"""
-        ...
-
     def _headers(self) -> dict:
         """Return headers required for API requests"""
         ...

chunkr_ai/api/task_response.py

@@ -3,10 +3,10 @@ from typing import TypeVar, Optional, Generic
 from pydantic import BaseModel, PrivateAttr
 import asyncio
 
-from .config import Configuration, OutputConfiguration, OutputResponse, Status
+from .configuration import Configuration, OutputConfiguration, OutputResponse, Status
 from .protocol import ChunkrClientProtocol
 from .misc import prepare_upload_data
-from .decorators import anywhere, require_task
+from .decorators import anywhere, require_task, retry_on_429
 
 T = TypeVar("T", bound="TaskResponse")
 
@@ -21,16 +21,22 @@ class TaskResponse(BaseModel, Generic[T]):
     status: Status
     task_id: str
     task_url: Optional[str] = None
+    include_chunks: bool = False
+    _base64_urls: bool = False
     _client: Optional[ChunkrClientProtocol] = PrivateAttr(default=None)
 
-    def with_client(self, client: ChunkrClientProtocol) -> T:
+    def with_client(self, client: ChunkrClientProtocol, include_chunks: bool = False, base64_urls: bool = False) -> T:
         self._client = client
+        self.include_chunks = include_chunks
+        self._base64_urls = base64_urls
         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 getattr(self._client, 'raise_on_failure', True):
+                raise ValueError(self.message)
+            return self
         if self.status not in ("Starting", "Processing"):
             return self
         return None
@@ -43,11 +49,12 @@ class TaskResponse(BaseModel, Generic[T]):
             )
             r.raise_for_status()
             return r.json()
-        except (ConnectionError, TimeoutError) as _:
-            print("Connection error while polling the task, retrying...")
+        except (ConnectionError, TimeoutError, OSError) as e:
+            print(f"Connection error while polling the task: {str(e)}, retrying...")
             await asyncio.sleep(0.5)
-        except Exception:
-            raise
+            return await self._poll_request() 
+        except Exception as e:
+            raise e
 
     @anywhere()
     async def poll(self) -> T:
@@ -62,6 +69,7 @@ class TaskResponse(BaseModel, Generic[T]):
 
     @anywhere()
     @require_task()
+    @retry_on_429()
     async def update(self, config: Configuration) -> T:
         """Update the task configuration."""
         f = await prepare_upload_data(None, config, self._client._client)

chunkr_ai/models.py

@@ -1,4 +1,4 @@
-from .api.config import (
+from .api.configuration import (
     BoundingBox,
     Chunk,
     ChunkProcessing,

{chunkr_ai-0.0.23.dist-info → chunkr_ai-0.0.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.23
+Version: 0.0.25
 Summary: Python client for Chunkr: open source document intelligence
 Author-email: Ishaan Kapoor <ishaan@lumina.sh>
 License: MIT License
@@ -35,6 +35,7 @@ Provides-Extra: test
 Requires-Dist: pytest>=7.0.0; extra == "test"
 Requires-Dist: pytest-xdist>=3.0.0; extra == "test"
 Requires-Dist: pytest-asyncio>=0.21.0; extra == "test"
+Requires-Dist: ruff>=0.9.3; extra == "test"
 
 # Chunkr Python Client
 

chunkr_ai-0.0.25.dist-info/RECORD

@@ -0,0 +1,16 @@
+chunkr_ai/__init__.py,sha256=6KpYv2lmD6S5z2kc9pqwuLP5VDHmOuu2qDZArUIhb1s,53
+chunkr_ai/models.py,sha256=tOI7ylkhyeFfCLMisk96EPsH4UEcjBx1Mcisxc_AYXI,757
+chunkr_ai/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai/api/auth.py,sha256=hlv0GiUmlsbFO1wLL9sslqOnsBSoBqkL_6Mk2SDvxgE,413
+chunkr_ai/api/chunkr.py,sha256=VnbuAPlWLqyf8xCCU_kpdybgjVPTwZLarDQoD3uozY0,3065
+chunkr_ai/api/chunkr_base.py,sha256=giW56fL7xxJphdOTpIH52dXxpNt7OdP8pNiPSqbNjGM,5835
+chunkr_ai/api/configuration.py,sha256=0wnrKlUIO7opvV963Gr_S8tlAjpo_IkNmbTi1_FwEug,3751
+chunkr_ai/api/decorators.py,sha256=HSq3vcxOeUJkaWaf7HOvCyg9dWkVo8cG5BrU-jhbhmc,4053
+chunkr_ai/api/misc.py,sha256=5PBI6pvOXr0x-3WieSKLrC8MA0iGPa-IG-5FEZ3vnr0,5724
+chunkr_ai/api/protocol.py,sha256=LjPrYSq52m1afIlAo0yVGXlGZxPRh8J6g7S4PAit3Zo,388
+chunkr_ai/api/task_response.py,sha256=CZIa3w5qPvSZDbDJ-LAtg7OOY91LsruemaXNyO2PymI,4256
+chunkr_ai-0.0.25.dist-info/LICENSE,sha256=w3R12yNDyZpMiy2lxy_hvNbsldC75ww79sF0u11rkho,1069
+chunkr_ai-0.0.25.dist-info/METADATA,sha256=CG1cO9YX7TpAHwBXgqLDgF9nwmVv30WLsWzfULx06W4,6996
+chunkr_ai-0.0.25.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+chunkr_ai-0.0.25.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
+chunkr_ai-0.0.25.dist-info/RECORD,,

chunkr_ai-0.0.23.dist-info/RECORD

@@ -1,17 +0,0 @@
-chunkr_ai/__init__.py,sha256=6KpYv2lmD6S5z2kc9pqwuLP5VDHmOuu2qDZArUIhb1s,53
-chunkr_ai/models.py,sha256=MK8FPbWDj1ynvSHaYuslKCPybxLuAlrsVIM3Eym3kKI,750
-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=XTXJFs0xjYY3w3N4fSQcxtJFBtNfzFYYkh6nDlFz4cY,2714
-chunkr_ai/api/chunkr_base.py,sha256=OkycHDHkdGX953_ab0XdYBnPDzSXYE30L3j52hBb8D0,5046
-chunkr_ai/api/config.py,sha256=NmPTsDvcjkvNx0gNzDTz-oFG5rQC7jm-H70O_crJCw8,4478
-chunkr_ai/api/decorators.py,sha256=y_Z9z0O2XXiX9z6jWDwdbCPdQyMLnjE0pGkJjHQEv_Q,2652
-chunkr_ai/api/misc.py,sha256=5Q2K713VPwf3S2519KTzjT9PKhTEBgBMk1d8NNnmpZ0,5717
-chunkr_ai/api/protocol.py,sha256=Nt8aWr4ouVwCvoLqVI5vnXJhT2cvxt0sQC-svUk2G5w,458
-chunkr_ai/api/task_response.py,sha256=aAx7otuvsp-A0U5EaHRkbnRJMoLI8N4lOMo8bS8emJc,3843
-chunkr_ai-0.0.23.dist-info/LICENSE,sha256=w3R12yNDyZpMiy2lxy_hvNbsldC75ww79sF0u11rkho,1069
-chunkr_ai-0.0.23.dist-info/METADATA,sha256=QsO__q1V9SJYz2uugyzj_CZpOuPie6AhLfY39hTNaOM,6952
-chunkr_ai-0.0.23.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-chunkr_ai-0.0.23.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
-chunkr_ai-0.0.23.dist-info/RECORD,,