chunkr-ai 0.0.17__py3-none-any.whl → 0.0.18__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.18.dist-info}/METADATA

@@ -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.18.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.18.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai-0.0.18.dist-info/METADATA,sha256=U9774MiX43ALUNaQRS6y05Kcg4q8REMVu5lyds1DJUo,4807
+chunkr_ai-0.0.18.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+chunkr_ai-0.0.18.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
+chunkr_ai-0.0.18.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()

chunkr_ai/api/schema.py

@@ -1,136 +0,0 @@
-from pydantic import BaseModel
-from typing import Optional, List, Union, Type
-import json
-
-
-class Property(BaseModel):
-    name: str
-    prop_type: str
-    description: Optional[str] = None
-    default: Optional[str] = None
-
-
-class JsonSchema(BaseModel):
-    title: str
-    properties: List[Property]
-
-
-def from_pydantic(
-    pydantic: Union[BaseModel, Type[BaseModel]], current_depth: int = 0
-) -> dict:
-    """Convert a Pydantic model to a Chunk json schema."""
-    MAX_DEPTH = 5
-    model = pydantic if isinstance(pydantic, type) else pydantic.__class__
-    schema = model.model_json_schema()
-    properties = []
-
-    def get_enum_description(details: dict) -> str:
-        """Get description including enum values if they exist"""
-        description = details.get("description", "")
-
-        # First check if this is a direct enum
-        if "enum" in details:
-            enum_values = details["enum"]
-            enum_str = "\nAllowed values:\n" + "\n".join(
-                f"- {val}" for val in enum_values
-            )
-            return f"{description}{enum_str}"
-
-        # Then check if it's a reference to an enum
-        if "$ref" in details:
-            ref_schema = resolve_ref(details["$ref"], schema.get("$defs", {}))
-            if "enum" in ref_schema:
-                enum_values = ref_schema["enum"]
-                enum_str = "\nAllowed values:\n" + "\n".join(
-                    f"- {val}" for val in enum_values
-                )
-                return f"{description}{enum_str}"
-
-        return description
-
-    def resolve_ref(ref: str, definitions: dict) -> dict:
-        """Resolve a $ref reference to its actual schema"""
-        if not ref.startswith("#/$defs/"):
-            return {}
-        ref_name = ref[len("#/$defs/") :]
-        return definitions.get(ref_name, {})
-
-    def get_nested_schema(field_schema: dict, depth: int) -> dict:
-        if depth >= MAX_DEPTH:
-            return {}
-
-        # If there's a $ref, resolve it first
-        if "$ref" in field_schema:
-            field_schema = resolve_ref(field_schema["$ref"], schema.get("$defs", {}))
-
-        nested_props = {}
-        if field_schema.get("type") == "object":
-            for name, details in field_schema.get("properties", {}).items():
-                if details.get("type") == "object" or "$ref" in details:
-                    ref_schema = details
-                    if "$ref" in details:
-                        ref_schema = resolve_ref(
-                            details["$ref"], schema.get("$defs", {})
-                        )
-                    nested_schema = get_nested_schema(ref_schema, depth + 1)
-                    nested_props[name] = {
-                        "type": "object",
-                        "description": get_enum_description(details),
-                        "properties": nested_schema,
-                    }
-                else:
-                    nested_props[name] = {
-                        "type": details.get("type", "string"),
-                        "description": get_enum_description(details),
-                    }
-        return nested_props
-
-    for name, details in schema.get("properties", {}).items():
-        # Handle arrays
-        if details.get("type") == "array":
-            items = details.get("items", {})
-            if "$ref" in items:
-                items = resolve_ref(items["$ref"], schema.get("$defs", {}))
-
-            # Get nested schema for array items
-            item_schema = get_nested_schema(items, current_depth)
-            description = get_enum_description(details)
-
-            if item_schema:
-                description = f"{description}\nList items schema:\n{json.dumps(item_schema, indent=2)}"
-
-            prop = Property(name=name, prop_type="list", description=description)
-        # Handle objects and references
-        elif details.get("type") == "object" or "$ref" in details:
-            prop_type = "object"
-            ref_schema = details
-            if "$ref" in details:
-                ref_schema = resolve_ref(details["$ref"], schema.get("$defs", {}))
-
-            nested_schema = get_nested_schema(ref_schema, current_depth)
-
-            prop = Property(
-                name=name,
-                prop_type=prop_type,
-                description=get_enum_description(details),
-                properties=nested_schema,
-            )
-
-        # Handle primitive types
-        else:
-            prop = Property(
-                name=name,
-                prop_type=details.get("type", "string"),
-                description=get_enum_description(details),
-                default=str(details.get("default"))
-                if details.get("default") is not None
-                else None,
-            )
-
-        properties.append(prop)
-
-    json_schema = JsonSchema(
-        title=schema.get("title", model.__name__), properties=properties
-    )
-
-    return json_schema.model_dump(mode="json", exclude_none=True)