chunkr-ai 0.0.11__py3-none-any.whl → 0.0.14__py3-none-any.whl

chunkr_ai/api/base.py

@@ -0,0 +1,173 @@
+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.py

@@ -15,86 +15,17 @@ class Chunkr(ChunkrBase):
         self._session = requests.Session()
 
     def upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> TaskResponse:
-        """Upload a file and wait for processing to complete.
-
-        Args:
-            file: The file to upload. 
-            config: Configuration options for processing. Optional.
-
-        Examples:
-        ```
-        # Upload from file path
-        chunkr.upload("document.pdf")
-
-        # Upload from URL
-        chunkr.upload("https://example.com/document.pdf")
-
-        # Upload from base64 string (must include MIME type header)
-        chunkr.upload("data:application/pdf;base64,JVBERi0xLjcKCjEgMCBvYmo...")
-
-        # Upload from opened file
-        with open("document.pdf", "rb") as f:
-            chunkr.upload(f)
-
-        # Upload an image
-        from PIL import Image
-        img = Image.open("photo.jpg")
-        chunkr.upload(img)
-        ```
-        Returns:
-            TaskResponse: The completed task response
-        """
         task = self.create_task(file, config)
         return task.poll()
     
     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
-        """
         task = self.update_task(task_id, config)
         return task.poll()
 
     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. 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 = chunkr.start_upload("document.pdf")
-
-        # Upload from opened file
-        with open("document.pdf", "rb") as f:
-            task = chunkr.start_upload(f)
-
-        # Upload from URL
-        task = chunkr.start_upload("https://example.com/document.pdf")
-
-        # Upload from base64 string (must include MIME type header)
-        task = chunkr.start_upload("data:application/pdf;base64,JVBERi0xLjcKCjEgMCBvYmo...")
-
-        # Upload an image
-        from PIL import Image
-        img = Image.open("photo.jpg")
-        task = chunkr.start_upload(img)
-
-        # Wait for the task to complete - this can be done when needed
-        task.poll()
-        ```
-
-        Returns:
-            TaskResponse: The initial task response
-        """
         files= prepare_upload_data(file, config)
+        if not self._session:
+            raise ValueError("Session not found")
         r = self._session.post(
             f"{self.url}/api/v1/task",
             files=files,
@@ -104,16 +35,9 @@ class Chunkr(ChunkrBase):
         return TaskResponse(**r.json()).with_client(self)
     
     def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
-        """Update a task by its ID.
-        
-        Args:
-            task_id: The ID of the task to update
-            config: The new configuration to use
-
-        Returns:
-            TaskResponse: The updated task response
-        """
         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,
@@ -124,14 +48,8 @@ class Chunkr(ChunkrBase):
         return TaskResponse(**r.json()).with_client(self)
     
     def get_task(self, task_id: str) -> TaskResponse:
-        """Get a task response by its ID.
-        
-        Args:
-            task_id: The ID of the task to get
-
-        Returns:
-            TaskResponse: The task response
-        """
+        if not self._session:
+            raise ValueError("Session not found")
         r = self._session.get(
             f"{self.url}/api/v1/task/{task_id}",
             headers=self._headers()
@@ -141,11 +59,8 @@ class Chunkr(ChunkrBase):
 
 
     def delete_task(self, task_id: str) -> None:
-        """Delete a task by its ID.
-        
-        Args:
-            task_id: The ID of the task to delete
-        """
+        if not self._session:
+            raise ValueError("Session not found")
         r = self._session.delete(
             f"{self.url}/api/v1/task/{task_id}",
             headers=self._headers()
@@ -153,11 +68,8 @@ class Chunkr(ChunkrBase):
         r.raise_for_status()
 
     def cancel_task(self, task_id: str) -> None:
-        """Cancel a task by its ID.
-        
-        Args:
-            task_id: The ID of the task to cancel
-        """
+        if not self._session:
+            raise ValueError("Session not found")
         r = self._session.get(
             f"{self.url}/api/v1/task/{task_id}/cancel",
             headers=self._headers()

chunkr_ai/api/chunkr_async.py

@@ -1,11 +1,11 @@
 from .chunkr_base import ChunkrBase
-from .task import TaskResponse
 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
-from .misc import prepare_upload_data
 
 class ChunkrAsync(ChunkrBase):
     """Asynchronous Chunkr API client"""
@@ -14,129 +14,99 @@ class ChunkrAsync(ChunkrBase):
         super().__init__(url, api_key)
         self._client = httpx.AsyncClient()
 
-    async def upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> TaskResponse:
-        """Upload a file and wait for processing to complete.
-
-        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,JVBERi0xLjcKCjEgMCBvYmo...")
-
-        # Upload an image
-        from PIL import Image
-        img = Image.open("photo.jpg")
-        await chunkr.upload(img)
-        ```
-        Returns:
-            TaskResponse: The completed task response
-        """
-        task = await self.create_task(file, config)
-        return await task.poll_async()
+    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) -> 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
-        """
-        task = await self.update_task(task_id, config)
-        return await task.poll_async()
-
-    async 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. It will not wait for processing to complete. To wait for the full processing to complete, use `task.poll_async()`.
-
-        Args:
-            file: The file to upload.
-            config: Configuration options for processing. Optional.
-
-        Examples:
-        ```
-        # Upload from file path
-        task = await chunkr.start_upload("document.pdf")
-
-        # Upload from opened file
-        with open("document.pdf", "rb") as f:
-            task = await chunkr.start_upload(f)
-    
-        # Upload from URL
-        task = await chunkr.start_upload("https://example.com/document.pdf")
-
-        # Upload from base64 string (must include MIME type header)
-        task = await chunkr.start_upload("data:application/pdf;base64,JVBERi0xLjcKCjEgMCBvYmo...")
-
-        # Upload an image
-        from PIL import Image
-        img = Image.open("photo.jpg")
-        task = await chunkr.start_upload(img)
-
-        # Wait for the task to complete - this can be done when needed
-        await task.poll_async()
-        ```
-
-        Returns:
-            TaskResponse: The initial task response
-        """
-        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)
-
-    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()
-        )
+    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 TaskResponse(**r.json()).with_client(self)
+            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) -> 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)
+    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:
-        r = await self._client.delete(
-            f"{self.url}/api/v1/task/{task_id}",
-            headers=self._headers()
-        )
-        r.raise_for_status()
+        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:
-        r = await self._client.get(
-            f"{self.url}/api/v1/task/{task_id}/cancel",
-            headers=self._headers()
-        )
-        r.raise_for_status()
+        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
 

chunkr_ai/api/chunkr_base.py

@@ -1,5 +1,6 @@
 from .config import Configuration
 from .task import TaskResponse
+from .task_async import TaskResponseAsync
 from .auth import HeadersMixin
 from abc import abstractmethod
 from dotenv import load_dotenv
@@ -28,42 +29,107 @@ class ChunkrBase(HeadersMixin):
         self.url = self.url.rstrip("/")
 
     @abstractmethod
-    def upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> TaskResponse:
+    def upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> Union[TaskResponse, TaskResponseAsync]:
         """Upload a file and wait for processing to complete.
+
+        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)
         
-        Must be implemented by subclasses.
+        # 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 update_task(self, task_id: str, config: Configuration) -> TaskResponse:
-        """Update a task by its ID.
+    def update(self, task_id: str, config: Configuration) -> Union[TaskResponse, TaskResponseAsync]:
+        """Update a task by its ID and wait for processing to complete.
         
-        Must be implemented by subclasses.
+        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.
-        
-        Must be implemented by subclasses.
+    def create_task(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> Union[TaskResponse, TaskResponseAsync]:
+        """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")
+
+        # 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.
+    def update_task(self, task_id: str, config: Configuration) -> Union[TaskResponse, TaskResponseAsync]:
+        """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()`.
         
-        Must be implemented by subclasses.
+        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 get_task(self, task_id: str) -> TaskResponse:
+    def get_task(self, task_id: str) -> Union[TaskResponse, TaskResponseAsync]:
         """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
 
@@ -71,7 +137,8 @@ class ChunkrBase(HeadersMixin):
     def delete_task(self, task_id: str) -> None:
         """Delete a task by its ID.
         
-        Must be implemented by subclasses.
+        Args:
+            task_id: The ID of the task to delete
         """
         pass
     
@@ -79,7 +146,8 @@ class ChunkrBase(HeadersMixin):
     def cancel_task(self, task_id: str) -> None:
         """Cancel a task by its ID.
         
-        Must be implemented by subclasses.
+        Args:
+            task_id: The ID of the task to cancel
         """
         pass
 

chunkr_ai/api/config.py

@@ -142,4 +142,4 @@ class Status(str, Enum):
     PROCESSING = "Processing"
     SUCCEEDED = "Succeeded"
     FAILED = "Failed"
-    CANCELLED = "Cancelled"
+    CANCELLED = "Cancelled"

chunkr_ai/api/protocol.py

@@ -1,14 +1,14 @@
-from typing import runtime_checkable, Protocol
+from typing import Optional, runtime_checkable, Protocol
 from requests import Session
-from httpx import AsyncClient
+from aiohttp import ClientSession
 
 @runtime_checkable
 class ChunkrClientProtocol(Protocol):
     """Protocol defining the interface for Chunkr clients"""
     url: str
     _api_key: str
-    _session: Session
-    _client: AsyncClient
+    _session: Optional[Session] = None
+    _client: Optional[ClientSession] = None
 
     def get_api_key(self) -> str:
         """Get the API key"""

chunkr_ai/api/task.py

@@ -1,40 +1,16 @@
-from .protocol import ChunkrClientProtocol
-from .config import Configuration, OutputResponse, Status
+from .config import Configuration
 from .misc import prepare_upload_data
-import asyncio
-from datetime import datetime
-from pydantic import BaseModel, PrivateAttr
+from .task_base import TaskBase
 import time
-from typing import Optional, Union
-
-class TaskResponse(BaseModel):
-    configuration: Configuration
-    created_at: datetime
-    expires_at: Optional[datetime] = None
-    file_name: Optional[str] = None
-    finished_at: Optional[datetime] = None
-    input_file_url: Optional[str] = None
-    message: str
-    output: Optional[OutputResponse] = None
-    page_count: Optional[int] = None
-    pdf_url: Optional[str] = None
-    started_at: Optional[datetime] = None
-    status: Status
-    task_id: str
-    task_url: Optional[str] = None
-    _client: Optional[Union[ChunkrClientProtocol]] = PrivateAttr(default=None)
-
-    def with_client(self, client: Union[ChunkrClientProtocol]) -> 'TaskResponse':
-        self._client = client
-        return self
-    
-    def _poll_request_sync(self) -> dict:
-        """Helper method to make polling request with retry logic (synchronous)"""
-        if not self.task_url:
-            raise ValueError("Task URL not found in response")
 
+class TaskResponse(TaskBase):
+    def _poll_request(self) -> dict:
         while True:
             try:
+                if not self.task_url:
+                   raise ValueError("Task URL not found in response")
+                if not self._client._session:
+                    raise ValueError("Client session not found")
                 r = self._client._session.get(self.task_url, headers=self._client._headers())
                 r.raise_for_status()
                 return r.json()
@@ -44,125 +20,50 @@ class TaskResponse(BaseModel):
             except Exception as e:
                 raise
 
-    async def _poll_request_async(self) -> dict:
-        """Helper method to make polling request with retry logic (asynchronous)"""
-        if not self.task_url:
-            raise ValueError("Task URL not found in response")
-
-        while True:
-            try:
-                r = await self._client._client.get(self.task_url, headers=self._client._headers())
-                r.raise_for_status()
-                response = r.json()
-                return response
-            except (ConnectionError, TimeoutError) as _:
-                print("Connection error while polling the task, retrying...")
-                await asyncio.sleep(0.5)
-            except Exception as e:
-                raise
-
-    def _check_status(self) -> Optional['TaskResponse']:
-        """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
-
     def poll(self) -> 'TaskResponse':
-        """Poll the task for completion."""
         while True:
-            response = self._poll_request_sync()
+            response = self._poll_request()
             updated_task = TaskResponse(**response).with_client(self._client)
             self.__dict__.update(updated_task.__dict__)
-            
             if result := self._check_status():
                 return result
-            
             time.sleep(0.5)
-
-    async def poll_async(self) -> 'TaskResponse':
-        """Poll the task for completion asynchronously."""
-        while True:
-            response = await self._poll_request_async()
-            updated_task = TaskResponse(**response).with_client(self._client)
-            self.__dict__.update(updated_task.__dict__)
-            
-            if result := self._check_status():
-                return result
-            
-            await asyncio.sleep(0.5)
-
-    def _get_content(self, content_type: str) -> str:
-        """Helper method to get either HTML, Markdown, or raw content."""
-        if not self.output:
-            return ""
-        parts = []
-        for c in self.output.chunks:
-            for s in c.segments:
-                content = getattr(s, content_type)
-                if content:
-                    parts.append(content)
-        return "\n".join(parts)
     
     def update(self, config: Configuration) -> 'TaskResponse':
+        if not self.task_url:
+            raise ValueError("Task URL not found")
+        if not self._client._session:
+            raise ValueError("Client session not found")
         files = prepare_upload_data(None, config)
         r = self._client._session.patch(
-            f"{self.task_url}",
+            self.task_url,
             files=files,
             headers=self._client._headers()
         )
         r.raise_for_status()
-        return TaskResponse(**r.json()).with_client(self._client)
-    
-    async def update_async(self, config: Configuration) -> 'TaskResponse':
-        files = prepare_upload_data(None, config)
-        r = await self._client._client.patch(
-            f"{self.task_url}",
-            files=files,
-            headers=self._client._headers()
-        )   
-        r.raise_for_status()
-        return TaskResponse(**r.json()).with_client(self._client)
+        updated = TaskResponse(**r.json()).with_client(self._client)
+        self.__dict__.update(updated.__dict__)
+        return self.poll()
     
     def cancel(self):
+        if not self.task_url:
+            raise ValueError("Task URL not found")
+        if not self._client._session:
+            raise ValueError("Client session not found")
         r = self._client._session.get(
             f"{self.task_url}/cancel",
             headers=self._client._headers()
         )
         r.raise_for_status()
         self.poll()
-    
-    async def cancel_async(self):
-        r = await self._client._client.get(
-            f"{self.task_url}/cancel",
-            headers=self._client._headers()
-        )
-        r.raise_for_status()
-        await self.poll_async()
 
     def delete(self):
+        if not self.task_url:
+            raise ValueError("Task URL not found")
+        if not self._client._session:
+            raise ValueError("Client session not found")
         r = self._client._session.delete(
-            f"{self.task_url}",
-            headers=self._client._headers()
-        )
-        r.raise_for_status()
-    
-    async def delete_async(self):
-        r = await self._client._client.delete(
-            f"{self.task_url}",
+            self.task_url,
             headers=self._client._headers()
         )
         r.raise_for_status()
-
-    def html(self) -> str:
-        """Get full HTML for the task"""
-        return self._get_content("html")
-
-    def markdown(self) -> str:
-        """Get full markdown for the task"""
-        return self._get_content("markdown")
-        
-    def content(self) -> str:
-        """Get full text for the task"""
-        return self._get_content("content")

chunkr_ai/api/task_async.py

@@ -1,34 +1,27 @@
-import asyncio
-from pydantic import BaseModel, PrivateAttr
-from datetime import datetime
-from typing import Optional, Union
-from .task_base import TaskBase
-from .protocol import ChunkrClientProtocol
-from .config import Configuration, OutputResponse, Status
+from .config import Configuration
 from .misc import prepare_upload_data
+from .task_base import TaskBase
+import asyncio
 
-class TaskResponseAsync(BaseModel, TaskBase):
-    configuration: Configuration
-    created_at: datetime
-    expires_at: Optional[datetime]
-    file_name: Optional[str]
-    finished_at: Optional[datetime]
-    input_file_url: Optional[str]
-    message: str
-    output: Optional[OutputResponse]
-    page_count: Optional[int]
-    pdf_url: Optional[str]
-    started_at: Optional[datetime]
-    status: Status
-    task_id: str
-    task_url: Optional[str]
-    _client: Optional[Union[ChunkrClientProtocol]] = PrivateAttr(default=None)
-
-    def with_client(self, client: Union[ChunkrClientProtocol]) -> 'TaskResponseAsync':
-        self._client = client
-        return self
+class TaskResponseAsync(TaskBase):
+    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 as e:
+            raise
 
     async def poll(self) -> 'TaskResponseAsync':
+        if not self.task_url:
+            raise ValueError("Task URL not found")
+        if not self._client._client:
+            raise ValueError("Client not found")
         while True:
             j = await self._poll_request()
             updated = TaskResponseAsync(**j).with_client(self._client)
@@ -37,31 +30,11 @@ class TaskResponseAsync(BaseModel, TaskBase):
                 return res
             await asyncio.sleep(0.5)
 
-    async def _poll_request(self) -> dict:
-        if not self.task_url:
-            raise ValueError("Task URL not found")
-        while True:
-            try:
-                r = await self._client._client.get(self.task_url, headers=self._client._headers())
-                r.raise_for_status()
-                return r.json()
-            except Exception as e:
-                if self.status == Status.FAILED:
-                    raise ValueError(self.message) from e
-                await asyncio.sleep(0.5)
-
-    def _check_status(self) -> Optional['TaskResponseAsync']:
-        if self.status == Status.FAILED:
-            raise ValueError(f"Task failed: {self.message}")
-        if self.status == Status.CANCELLED:
-            return self
-        if self.status not in [Status.STARTING, Status.PROCESSING]:
-            return self
-        return None
-
     async def update(self, config: Configuration) -> 'TaskResponseAsync':
         if not self.task_url:
             raise ValueError("Task URL not found")
+        if not self._client._client:
+            raise ValueError("Client not found")
         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()
@@ -72,32 +45,16 @@ class TaskResponseAsync(BaseModel, TaskBase):
     async def cancel(self):
         if not self.task_url:
             raise ValueError("Task URL not found")
+        if not self._client._client:
+            raise ValueError("Client not found")
         r = await self._client._client.get(f"{self.task_url}/cancel", headers=self._client._headers())
         r.raise_for_status()
         return await self.poll()
 
     async def delete(self):
+        if not self.task_url:
+            raise ValueError("Task URL not found")
+        if not self._client._client:
+            raise ValueError("Client not found")
         r = await self._client._client.delete(self.task_url, headers=self._client._headers())
         r.raise_for_status()
-
-    def html(self) -> str:
-        return self._get_content("html")
-
-    def markdown(self) -> str:
-        return self._get_content("markdown")
-
-    def content(self) -> str:
-        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)
-
-    # Satisfying TaskBase abstract methods with stubs

chunkr_ai/api/task_base.py

@@ -1,31 +1,82 @@
+from .config import Configuration, Status, OutputResponse
+from .protocol import ChunkrClientProtocol
 from abc import ABC, abstractmethod
-from .config import Configuration
+from typing import TypeVar, Optional, Generic
+from pydantic import BaseModel, PrivateAttr
+from datetime import datetime
+
+T = TypeVar('T', bound='TaskBase')
+
+class TaskBase(BaseModel, ABC, Generic[T]):
+    configuration: Configuration
+    created_at: datetime
+    expires_at: Optional[datetime]
+    file_name: Optional[str]
+    finished_at: Optional[datetime]
+    input_file_url: Optional[str]
+    message: str
+    output: Optional[OutputResponse]
+    page_count: Optional[int]
+    pdf_url: Optional[str]
+    started_at: Optional[datetime]
+    status: Status
+    task_id: str
+    task_url: Optional[str]
+    _client: Optional[ChunkrClientProtocol] = PrivateAttr(default=None)
 
-class TaskBase(ABC):
     @abstractmethod
-    def poll(self):
+    def _poll_request(self) -> dict:
+        """Helper method to make polling request with retry logic (synchronous)"""
         pass
 
     @abstractmethod
-    def update(self, config: Configuration):
+    def poll(self) -> T:
+        """Poll the task for completion."""
         pass
 
     @abstractmethod
-    def cancel(self):
+    def update(self, config: Configuration) -> T:
+        """Update the task configuration."""
         pass
 
     @abstractmethod
-    def delete(self):
+    def cancel(self) -> T:
+        """Cancel the task."""
         pass
 
     @abstractmethod
-    def html(self) -> str:
+    def delete(self) -> T:
+        """Delete the task."""
         pass
 
-    @abstractmethod
+    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
+
+    def html(self) -> str:
+        return self._get_content("html")
+
     def markdown(self) -> str:
-        pass
+        return self._get_content("markdown")
 
-    @abstractmethod
     def content(self) -> str:
-        pass
+        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-0.0.11.dist-info → chunkr_ai-0.0.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.11
+Version: 0.0.14
 Summary: Python client for Chunkr: open source document intelligence
 Author-email: Ishaan Kapoor <ishaan@lumina.sh>
 Project-URL: Homepage, https://chunkr.ai
@@ -80,7 +80,7 @@ async def process_document():
     # If you want to upload without waiting for processing
     task = await chunkr.start_upload("document.pdf")
     # ... do other things ...
-    await task.poll_async()  # Check status when needed
+    await task.poll()  # Check status when needed
 ```
 
 ### Additional Features

chunkr_ai-0.0.14.dist-info/RECORD

@@ -0,0 +1,21 @@
+chunkr_ai/__init__.py,sha256=eXygrEhGxxIHXNYIlHF2eied8rGsx2RphgR8Wo4lRyo,110
+chunkr_ai/models.py,sha256=-dbwtTHTcGhH3LXUdVUPkobbPoeFNXRizeAW8BCGSkE,903
+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=iSd5Jek2BFaHGw9HY-RrqgwP56BHFU0xbSuJS4fU6AA,425
+chunkr_ai/api/base.py,sha256=WDHx8tU0fl9_-yvYTKL-U0uaxHv-8_bRfiw9Xkl-mWM,6499
+chunkr_ai/api/chunkr.py,sha256=A1KEjA4zRJkGZhYwhWde4CFncrljEMky4SO4LPzYvws,2652
+chunkr_ai/api/chunkr_async.py,sha256=OvXd9Ma3rsp6q8nJsWzNgTKAGk-py93DqAENE8fMQfc,4153
+chunkr_ai/api/chunkr_base.py,sha256=run4UJVKa7Gx8I_ME0Mol-c_b-NIcexNkgxHu_hbX5M,4996
+chunkr_ai/api/config.py,sha256=joTn7jiOlJXTwwza-jHauLV-39CMzaxZVGB9JBm8Cok,4862
+chunkr_ai/api/misc.py,sha256=9vnfrbJ7sFlZqwEIQ4NTMb5rhPOmETT7e1jR-b42PXM,4977
+chunkr_ai/api/protocol.py,sha256=li-zy7Z-ChR9kZqJlixQv1kUYrmesPHxwUtnE5p16tQ,529
+chunkr_ai/api/schema.py,sha256=OeLOhBRXeRBgEImg0Q6O9Z10ojT6aSEVvwnDR8UeENo,4971
+chunkr_ai/api/task.py,sha256=j-Odecnbj3NjAGDyUNmZHgqsRDYhHs3xBWgsT2rrHjs,2517
+chunkr_ai/api/task_async.py,sha256=HpuVW928s-V3jjPa8L5i86lvKcpKbMhkCuV1nKSyXVA,2437
+chunkr_ai/api/task_base.py,sha256=9S8UCsrEOAH48PmOpLlyAifKDrtaUtrBMUmMZq4Dceg,2328
+chunkr_ai-0.0.14.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+chunkr_ai-0.0.14.dist-info/METADATA,sha256=1reEDprWgXg6bKCGKv6LXJxZZR-dQeJn9DeXGbp4Iwk,4839
+chunkr_ai-0.0.14.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+chunkr_ai-0.0.14.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
+chunkr_ai-0.0.14.dist-info/RECORD,,

chunkr_ai-0.0.11.dist-info/RECORD

@@ -1,19 +0,0 @@
-chunkr_ai/__init__.py,sha256=eXygrEhGxxIHXNYIlHF2eied8rGsx2RphgR8Wo4lRyo,110
-chunkr_ai/models.py,sha256=-dbwtTHTcGhH3LXUdVUPkobbPoeFNXRizeAW8BCGSkE,903
-chunkr_ai/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-chunkr_ai/api/auth.py,sha256=iSd5Jek2BFaHGw9HY-RrqgwP56BHFU0xbSuJS4fU6AA,425
-chunkr_ai/api/chunkr.py,sha256=0qpV9b1hOpDhA9EuKkXW9X_laUmw5NY3ZYq0cUOTbww,5190
-chunkr_ai/api/chunkr_async.py,sha256=ZkLBrn4cqzu3sqMfS8cfZZgSvpdyQuWZP95lfGxuHx0,4900
-chunkr_ai/api/chunkr_base.py,sha256=IYO0pmoL02GchIggj6_Q5nvtAUoOvYAAvT7VLFU6scY,2506
-chunkr_ai/api/config.py,sha256=y6wZz01ihRJ_5_cK_JklFWn397yll7jfXntd8bBBa5s,4861
-chunkr_ai/api/misc.py,sha256=9vnfrbJ7sFlZqwEIQ4NTMb5rhPOmETT7e1jR-b42PXM,4977
-chunkr_ai/api/protocol.py,sha256=XKS9RmtvBpJItYhPg18qlOCKpaSHdOuQTRSUxAdUz2g,479
-chunkr_ai/api/schema.py,sha256=OeLOhBRXeRBgEImg0Q6O9Z10ojT6aSEVvwnDR8UeENo,4971
-chunkr_ai/api/task.py,sha256=Z5Da_Ijvih5rBz5ry98oAYNcJEDbQhhDWBQ35nHCRK4,5881
-chunkr_ai/api/task_async.py,sha256=o7tXvViIrdcrdclxaGzxrgIv-n-W8-twQ7XsDLXfXhM,3659
-chunkr_ai/api/task_base.py,sha256=Tkk7dhIeB3ic5M9g_b-MVRdNv4XQTvajpaUy8JylQ8A,526
-chunkr_ai-0.0.11.dist-info/LICENSE,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-chunkr_ai-0.0.11.dist-info/METADATA,sha256=s8UeXDnBDVG_1RN5colcJCGhwrICRy9VMQWmTUKVRJc,4845
-chunkr_ai-0.0.11.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-chunkr_ai-0.0.11.dist-info/top_level.txt,sha256=0IZY7PZIiS8bw5r4NUQRUQ-ATi-L_3vLQVq3ZLouOW8,10
-chunkr_ai-0.0.11.dist-info/RECORD,,