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

chunkr_ai/__init__.py

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

chunkr_ai/api/chunkr.py

@@ -1,78 +1,85 @@
-from .chunkr_base import ChunkrBase
-from .config import Configuration
-from .task import TaskResponse
 from pathlib import Path
 from PIL import Image
-import requests
 from typing import Union, BinaryIO
-from .misc import prepare_upload_data
 
+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"""
-
-    def __init__(self, url: str = None, api_key: str = None):
-        super().__init__(url, api_key)
-        self._session = requests.Session()
-
-    def upload(
+    """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 = self.create_task(file, config)
-        return task.poll()
+        task = await self.create_task(file, config)
+        return await task.poll()
 
-    def update(self, task_id: str, config: Configuration) -> TaskResponse:
-        task = self.update_task(task_id, config)
-        return 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()
 
-    def create_task(
+    @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)
-        if not self._session:
-            raise ValueError("Session not found")
-        r = self._session.post(
+        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)
 
-    def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
+    @anywhere()
+    @ensure_client()
+    async def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
         files = prepare_upload_data(None, config)
-        if not self._session:
-            raise ValueError("Session not found")
-        r = self._session.patch(
-            f"{self.url}/api/v1/task/{task_id}", files=files, headers=self._headers()
+        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)
 
-    def get_task(self, task_id: str) -> TaskResponse:
-        if not self._session:
-            raise ValueError("Session not found")
-        r = self._session.get(
+    @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)
 
-    def delete_task(self, task_id: str) -> None:
-        if not self._session:
-            raise ValueError("Session not found")
-        r = self._session.delete(
+    @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()
 
-    def cancel_task(self, task_id: str) -> None:
-        if not self._session:
-            raise ValueError("Session not found")
-        r = self._session.get(
+    @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/api/chunkr_base.py

@@ -1,13 +1,16 @@
 from .config import Configuration
-from .task import TaskResponse
-from .task_async import TaskResponseAsync
+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
 from pathlib import Path
 from PIL import Image
-from typing import BinaryIO, Union
+import requests
+from typing import BinaryIO, Tuple, Union
 
 
 class ChunkrBase(HeadersMixin):
@@ -23,13 +26,138 @@ class ChunkrBase(HeadersMixin):
             )
 
         self.url = self.url.rstrip("/")
+        self._client = httpx.AsyncClient()
+
+    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,
-    ) -> Union[TaskResponse, TaskResponseAsync]:
+    ) -> TaskResponse:
         """Upload a file and wait for processing to complete.
 
         Args:
@@ -64,7 +192,7 @@ class ChunkrBase(HeadersMixin):
     @abstractmethod
     def update(
         self, task_id: str, config: Configuration
-    ) -> Union[TaskResponse, TaskResponseAsync]:
+    ) -> TaskResponse:
         """Update a task by its ID and wait for processing to complete.
 
         Args:
@@ -81,7 +209,7 @@ class ChunkrBase(HeadersMixin):
         self,
         file: Union[str, Path, BinaryIO, Image.Image],
         config: Configuration = None,
-    ) -> Union[TaskResponse, TaskResponseAsync]:
+    ) -> TaskResponse:
         """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:
@@ -117,7 +245,7 @@ class ChunkrBase(HeadersMixin):
     @abstractmethod
     def update_task(
         self, task_id: str, config: Configuration
-    ) -> Union[TaskResponse, TaskResponseAsync]:
+    ) -> 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:
@@ -130,7 +258,7 @@ class ChunkrBase(HeadersMixin):
         pass
 
     @abstractmethod
-    def get_task(self, task_id: str) -> Union[TaskResponse, TaskResponseAsync]:
+    def get_task(self, task_id: str) -> TaskResponse:
         """Get a task response by its ID.
 
         Args:
@@ -158,3 +286,9 @@ class ChunkrBase(HeadersMixin):
             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/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/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/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/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: