chunkr-ai 0.0.17__tar.gz → 0.0.18__tar.gz

{chunkr_ai-0.0.17/src/chunkr_ai.egg-info → chunkr_ai-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.17
+Version: 0.0.18
 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"

{chunkr_ai-0.0.17 → chunkr_ai-0.0.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chunkr-ai"
-version = "0.0.17"
+version = "0.0.18"
 authors = [{"name" = "Ishaan Kapoor", "email" = "ishaan@lumina.sh"}]
 description = "Python client for Chunkr: open source document intelligence"
 readme = "README.md"
@@ -16,7 +16,6 @@ dependencies = [
     "pydantic>=2.0.0",
     "pytest-asyncio>=0.21.0",
     "python-dotenv>=0.19.0",
-    "requests>=2.28.0",
 ]
 
 [project.optional-dependencies]

chunkr_ai-0.0.18/src/chunkr_ai/__init__.py

@@ -0,0 +1,3 @@
+from .api.chunkr import Chunkr
+
+__all__ = ["Chunkr"]

chunkr_ai-0.0.18/src/chunkr_ai/api/chunkr.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+from PIL import Image
+from typing import Union, BinaryIO
+
+from .config import Configuration
+from .decorators import anywhere, ensure_client
+from .misc import prepare_upload_data
+from .task_response import TaskResponse
+from .chunkr_base import ChunkrBase
+
+class Chunkr(ChunkrBase):
+    """Chunkr API client that works in both sync and async contexts"""
+    
+    @anywhere()
+    @ensure_client()
+    async def upload(
+        self,
+        file: Union[str, Path, BinaryIO, Image.Image],
+        config: Configuration = None,
+    ) -> TaskResponse:
+        task = await self.create_task(file, config)
+        return await task.poll()
+
+    @anywhere()
+    @ensure_client()
+    async def update(self, task_id: str, config: Configuration) -> TaskResponse:
+        task = await self.update_task(task_id, config)
+        return await task.poll()
+
+    @anywhere()
+    @ensure_client()
+    async def create_task(
+        self,
+        file: Union[str, Path, BinaryIO, Image.Image],
+        config: Configuration = None,
+    ) -> TaskResponse:
+        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 TaskResponse(**r.json()).with_client(self)
+
+    @anywhere()
+    @ensure_client()
+    async def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
+        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 TaskResponse(**r.json()).with_client(self)
+
+    @anywhere()
+    @ensure_client()
+    async def get_task(self, task_id: str) -> TaskResponse:
+        r = await self._client.get(
+            f"{self.url}/api/v1/task/{task_id}", headers=self._headers()
+        )
+        r.raise_for_status()
+        return TaskResponse(**r.json()).with_client(self)
+
+    @anywhere()
+    @ensure_client()
+    async def delete_task(self, task_id: str) -> None:
+        r = await self._client.delete(
+            f"{self.url}/api/v1/task/{task_id}", headers=self._headers()
+        )
+        r.raise_for_status()
+
+    @ensure_client()
+    @anywhere()
+    async def cancel_task(self, task_id: str) -> None:
+        r = await self._client.get(
+            f"{self.url}/api/v1/task/{task_id}/cancel", headers=self._headers()
+        )
+        r.raise_for_status()
+
+    @anywhere()
+    async def close(self):
+        if self._client:
+            await self._client.aclose()
+            self._client = None

chunkr_ai-0.0.17/src/chunkr_ai/api/base.py → chunkr_ai-0.0.18/src/chunkr_ai/api/chunkr_base.py

@@ -1,8 +1,9 @@
 from .config import Configuration
-from .task import TaskResponse
+from .task_response import TaskResponse
 from .auth import HeadersMixin
 from abc import abstractmethod
 from dotenv import load_dotenv
+import httpx
 import io
 import json
 import os
@@ -25,6 +26,7 @@ class ChunkrBase(HeadersMixin):
             )
 
         self.url = self.url.rstrip("/")
+        self._client = httpx.AsyncClient()
 
     def _prepare_file(
         self, file: Union[str, Path, BinaryIO, Image.Image]
@@ -158,19 +160,100 @@ class ChunkrBase(HeadersMixin):
     ) -> TaskResponse:
         """Upload a file and wait for processing to complete.
 
-        Must be implemented by subclasses.
+        Args:
+            file: The file to upload.
+            config: Configuration options for processing. Optional.
+
+        Examples:
+        ```python
+        # Upload from file path
+        await chunkr.upload("document.pdf")
+
+        # Upload from opened file
+        with open("document.pdf", "rb") as f:
+            await chunkr.upload(f)
+
+        # Upload from URL
+        await chunkr.upload("https://example.com/document.pdf")
+
+        # Upload from base64 string (must include MIME type header)
+        await chunkr.upload("data:application/pdf;base64,JVBERi0...")
+
+        # Upload an image
+        from PIL import Image
+        img = Image.open("photo.jpg")
+        await chunkr.upload(img)
+        ```
+        Returns:
+            TaskResponse: The completed task response
         """
         pass
 
     @abstractmethod
-    def start_upload(
+    def update(
+        self, task_id: str, config: Configuration
+    ) -> TaskResponse:
+        """Update a task by its ID and wait for processing to complete.
+
+        Args:
+            task_id: The ID of the task to update
+            config: Configuration options for processing. Optional.
+
+        Returns:
+            TaskResponse: The updated task response
+        """
+        pass
+
+    @abstractmethod
+    def create_task(
         self,
         file: Union[str, Path, BinaryIO, Image.Image],
         config: Configuration = None,
     ) -> TaskResponse:
-        """Upload a file for processing and immediately return the task response.
+        """Upload a file for processing and immediately return the task response. It will not wait for processing to complete. To wait for the full processing to complete, use `task.poll()`.
+
+        Args:
+            file: The file to upload.
+            config: Configuration options for processing. Optional.
+
+        Examples:
+        ```
+        # Upload from file path
+        task = await chunkr.create_task("document.pdf")
 
-        Must be implemented by subclasses.
+        # Upload from opened file
+        with open("document.pdf", "rb") as f:
+            task = await chunkr.create_task(f)
+
+        # Upload from URL
+        task = await chunkr.create_task("https://example.com/document.pdf")
+
+        # Upload from base64 string (must include MIME type header)
+        task = await chunkr.create_task("data:application/pdf;base64,JVBERi0xLjcKCjEgMCBvYmo...")
+
+        # Upload an image
+        from PIL import Image
+        img = Image.open("photo.jpg")
+        task = await chunkr.create_task(img)
+
+        # Wait for the task to complete - this can be done when needed
+        await task.poll()
+        ```
+        """
+        pass
+
+    @abstractmethod
+    def update_task(
+        self, task_id: str, config: Configuration
+    ) -> TaskResponse:
+        """Update a task by its ID and immediately return the task response. It will not wait for processing to complete. To wait for the full processing to complete, use `task.poll()`.
+
+        Args:
+            task_id: The ID of the task to update
+            config: Configuration options for processing. Optional.
+
+        Returns:
+            TaskResponse: The updated task response
         """
         pass
 
@@ -178,6 +261,34 @@ class ChunkrBase(HeadersMixin):
     def get_task(self, task_id: str) -> TaskResponse:
         """Get a task response by its ID.
 
-        Must be implemented by subclasses.
+        Args:
+            task_id: The ID of the task to get
+
+        Returns:
+            TaskResponse: The task response
+        """
+        pass
+
+    @abstractmethod
+    def delete_task(self, task_id: str) -> None:
+        """Delete a task by its ID.
+
+        Args:
+            task_id: The ID of the task to delete
+        """
+        pass
+
+    @abstractmethod
+    def cancel_task(self, task_id: str) -> None:
+        """Cancel a task by its ID.
+
+        Args:
+            task_id: The ID of the task to cancel
         """
         pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the client connection. 
+        This should be called when you're done using the client to properly clean up resources."""
+        pass

{chunkr_ai-0.0.17 → chunkr_ai-0.0.18}/src/chunkr_ai/api/config.py

@@ -1,26 +1,21 @@
 from pydantic import BaseModel, Field, model_validator, ConfigDict
 from enum import Enum
 from typing import Optional, List, Dict, Union, Type
-from .schema import from_pydantic
-
 
 class GenerationStrategy(str, Enum):
     LLM = "LLM"
     AUTO = "Auto"
 
-
 class CroppingStrategy(str, Enum):
     ALL = "All"
     AUTO = "Auto"
 
-
 class GenerationConfig(BaseModel):
     html: Optional[GenerationStrategy] = None
     llm: Optional[str] = None
     markdown: Optional[GenerationStrategy] = None
     crop_image: Optional[CroppingStrategy] = None
 
-
 class SegmentProcessing(BaseModel):
     model_config = ConfigDict(populate_by_name=True, alias_generator=str.title)
 
@@ -39,46 +34,38 @@ class SegmentProcessing(BaseModel):
     page_footer: Optional[GenerationConfig] = Field(default=None, alias="PageFooter")
     page: Optional[GenerationConfig] = Field(default=None, alias="Page")
 
-
 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"
 
-
 class SegmentationStrategy(str, Enum):
     LAYOUT_ANALYSIS = "LayoutAnalysis"
     PAGE = "Page"
 
-
 class BoundingBox(BaseModel):
     left: float
     top: float
     width: float
     height: float
 
-
 class OCRResult(BaseModel):
     bbox: BoundingBox
     text: str
     confidence: Optional[float]
 
-
 class SegmentType(str, Enum):
     CAPTION = "Caption"
     FOOTNOTE = "Footnote"
@@ -93,7 +80,6 @@ class SegmentType(str, Enum):
     TEXT = "Text"
     TITLE = "Title"
 
-
 class Segment(BaseModel):
     bbox: BoundingBox
     content: str
@@ -107,42 +93,41 @@ class Segment(BaseModel):
     segment_id: str
     segment_type: SegmentType
 
-
 class Chunk(BaseModel):
     chunk_id: str
     chunk_length: int
     segments: List[Segment]
 
-
 class ExtractedJson(BaseModel):
     data: Dict
 
-
 class OutputResponse(BaseModel):
     chunks: List[Chunk]
-    extracted_json: Optional[ExtractedJson] = Field(default=None)
-
+    file_name: Optional[str]
+    page_count: Optional[int]
+    pdf_url: Optional[str]
 
 class Model(str, Enum):
     FAST = "Fast"
     HIGH_QUALITY = "HighQuality"
 
-class PipelineType(str, Enum):
+class Pipeline(str, Enum):
     AZURE = "Azure"
 
 class Configuration(BaseModel):
-    chunk_processing: Optional[ChunkProcessing] = Field(default=None)
-    expires_in: Optional[int] = Field(default=None)
-    high_resolution: Optional[bool] = Field(default=None)
-    json_schema: Optional[Union[JsonSchema, Type[BaseModel], BaseModel]] = Field(
-        default=None
-    )
-    model: Optional[Model] = Field(default=None)
-    ocr_strategy: Optional[OcrStrategy] = Field(default=None)
-    segment_processing: Optional[SegmentProcessing] = Field(default=None)
-    segmentation_strategy: Optional[SegmentationStrategy] = Field(default=None)
-    pipeline: Optional[PipelineType] = Field(default=None)
-
+    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
+    pipeline: Optional[Pipeline] = None
+    
+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:
@@ -151,19 +136,7 @@ class Configuration(BaseModel):
                 values["chunk_processing"] = values.get("chunk_processing", {}) or {}
                 values["chunk_processing"]["target_length"] = target_length
         return values
-
-    @model_validator(mode="after")
-    def convert_json_schema(self) -> "Configuration":
-        if self.json_schema is not None and not isinstance(
-            self.json_schema, JsonSchema
-        ):
-            if isinstance(self.json_schema, (BaseModel, type)) and issubclass(
-                getattr(self.json_schema, "__class__", type), BaseModel
-            ):
-                self.json_schema = JsonSchema(**from_pydantic(self.json_schema))
-        return self
-
-
+    
 class Status(str, Enum):
     STARTING = "Starting"
     PROCESSING = "Processing"

chunkr_ai-0.0.18/src/chunkr_ai/api/decorators.py

@@ -0,0 +1,58 @@
+import functools
+import asyncio
+import httpx
+from typing import Callable, Any, TypeVar, Awaitable, ParamSpec, Union, overload
+
+T = TypeVar('T')
+P = ParamSpec('P')
+
+_sync_loop = None
+
+@overload
+def anywhere() -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Union[Awaitable[T], T]]]: ...
+
+def anywhere():
+    """Decorator that allows an async function to run anywhere - sync or async context."""
+    def decorator(async_func: Callable[P, Awaitable[T]]) -> Callable[P, Union[Awaitable[T], T]]:
+        @functools.wraps(async_func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Union[Awaitable[T], T]:
+            global _sync_loop
+            try:
+                loop = asyncio.get_running_loop()
+                return async_func(*args, **kwargs)
+            except RuntimeError:
+                if _sync_loop is None:
+                    _sync_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(_sync_loop)
+                try:
+                    return _sync_loop.run_until_complete(async_func(*args, **kwargs))
+                finally:
+                    asyncio.set_event_loop(None)
+        return wrapper
+    return decorator
+
+def ensure_client() -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+    """Decorator that ensures a valid httpx.AsyncClient exists before executing the method"""
+    def decorator(async_func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+        @functools.wraps(async_func)
+        async def wrapper(self: Any, *args: P.args, **kwargs: P.kwargs) -> T:
+            if not self._client or self._client.is_closed:
+                self._client = httpx.AsyncClient()
+            return await async_func(self, *args, **kwargs)
+        return wrapper
+    return decorator
+
+def require_task() -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+    """Decorator that ensures task has required attributes and valid client before execution"""
+    def decorator(async_func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+        @functools.wraps(async_func)
+        async def wrapper(self: Any, *args: P.args, **kwargs: P.kwargs) -> T:
+            if not self.task_url:
+                raise ValueError("Task URL not found")
+            if not self._client:
+                raise ValueError("Client not found")
+            if not self._client._client or self._client._client.is_closed:
+                self._client._client = httpx.AsyncClient()
+            return await async_func(self, *args, **kwargs)
+        return wrapper
+    return decorator

{chunkr_ai-0.0.17 → chunkr_ai-0.0.18}/src/chunkr_ai/api/misc.py

@@ -3,11 +3,9 @@ import io
 import json
 from pathlib import Path
 from PIL import Image
-from pydantic import BaseModel
 import requests
 from typing import Union, Tuple, BinaryIO, Optional
 
-
 def prepare_file(file: Union[str, Path, BinaryIO, Image.Image]) -> Tuple[str, BinaryIO]:
     """Convert various file types into a tuple of (filename, file-like object)."""
     # Handle URLs

{chunkr_ai-0.0.17 → chunkr_ai-0.0.18}/src/chunkr_ai/api/protocol.py

@@ -1,5 +1,4 @@
 from typing import Optional, runtime_checkable, Protocol
-from requests import Session
 from httpx import AsyncClient
 
 
@@ -9,7 +8,6 @@ class ChunkrClientProtocol(Protocol):
 
     url: str
     _api_key: str
-    _session: Optional[Session] = None
     _client: Optional[AsyncClient] = None
 
     def get_api_key(self) -> str:

chunkr_ai-0.0.18/src/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)