chunkr-ai 0.0.6__tar.gz → 0.0.8__tar.gz

{chunkr_ai-0.0.6/src/chunkr_ai.egg-info → chunkr_ai-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.6
+Version: 0.0.8
 Summary: Python client for Chunkr: open source document intelligence
 Author-email: Ishaan Kapoor <ishaan@lumina.sh>
 Project-URL: Homepage, https://chunkr.ai

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chunkr-ai"
-version = "0.0.6"
+version = "0.0.8"
 authors = [{"name" = "Ishaan Kapoor", "email" = "ishaan@lumina.sh"}]
 description = "Python client for Chunkr: open source document intelligence"
 readme = "README.md"

chunkr_ai-0.0.8/src/chunkr_ai/api/base.py

@@ -0,0 +1,85 @@
+from .config import Configuration
+from .task import TaskResponse
+from .auth import HeadersMixin
+from abc import abstractmethod
+from dotenv import load_dotenv
+import os
+from pathlib import Path
+from PIL import Image
+from typing import BinaryIO, Union
+
+class ChunkrBase(HeadersMixin):
+    """Base class with shared functionality for Chunkr API clients."""
+
+    def __init__(self, url: str = None, api_key: str = None):
+        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("/")
+
+    @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 update_task(self, task_id: str, config: Configuration) -> TaskResponse:
+        """Update a task by its ID.
+        
+        Must be implemented by subclasses.
+        """
+        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.
+        """
+        pass
+
+    @abstractmethod
+    def update_task(self, task_id: str, config: Configuration) -> TaskResponse:
+        """Update a task by its ID.
+        
+        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
+
+    @abstractmethod
+    def delete_task(self, task_id: str) -> None:
+        """Delete a task by its ID.
+        
+        Must be implemented by subclasses.
+        """
+        pass
+    
+    @abstractmethod
+    def cancel_task(self, task_id: str) -> None:
+        """Cancel a task by its ID.
+        
+        Must be implemented by subclasses.
+        """
+        pass
+

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai/api/chunkr.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from PIL import Image
 import requests
 from typing import Union, BinaryIO
+from .misc import prepare_upload_data
 
 class Chunkr(ChunkrBase):
     """Chunkr API client"""
@@ -43,10 +44,23 @@ class Chunkr(ChunkrBase):
         Returns:
             TaskResponse: The completed task response
         """
-        task = self.start_upload(file, config)
+        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.
 
-    def start_upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> TaskResponse:
+        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:
@@ -80,16 +94,35 @@ class Chunkr(ChunkrBase):
         Returns:
             TaskResponse: The initial task response
         """
-        files, data = self._prepare_upload_data(file, config)
+        files= prepare_upload_data(file, config)
         r = self._session.post(
             f"{self.url}/api/v1/task",
             files=files,
-            data=data,  
             headers=self._headers()
         )
         r.raise_for_status()
         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)
+        r = self._session.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:
         """Get a task response by its ID.
         
@@ -106,3 +139,29 @@ class Chunkr(ChunkrBase):
         r.raise_for_status()
         return TaskResponse(**r.json()).with_client(self)
 
+
+    def delete_task(self, task_id: str) -> None:
+        """Delete a task by its ID.
+        
+        Args:
+            task_id: The ID of the task to delete
+        """
+        r = self._session.delete(
+            f"{self.url}/api/v1/task/{task_id}",
+            headers=self._headers()
+        )
+        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
+        """
+        r = self._session.get(
+            f"{self.url}/api/v1/task/{task_id}/cancel",
+            headers=self._headers()
+        )
+        r.raise_for_status()
+
+  

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai/api/chunkr_async.py

@@ -5,6 +5,7 @@ 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"""
@@ -43,10 +44,23 @@ class ChunkrAsync(ChunkrBase):
         Returns:
             TaskResponse: The completed task response
         """
-        task = await self.start_upload(file, config)
+        task = await self.create_task(file, config)
         return await task.poll_async()
+    
+    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.
 
-    async def start_upload(self, file: Union[str, Path, BinaryIO, Image.Image], config: Configuration = None) -> TaskResponse:
+        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:
@@ -80,16 +94,26 @@ class ChunkrAsync(ChunkrBase):
         Returns:
             TaskResponse: The initial task response
         """
-        files, data = self._prepare_upload_data(file, config)
+        files = prepare_upload_data(file, config)
         r = await self._client.post(
             f"{self.url}/api/v1/task",
             files=files,
-            json=config.model_dump() if config else {},
             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()
+        )
+     
+        r.raise_for_status()
+        return TaskResponse(**r.json()).with_client(self)
+    
     async def get_task(self, task_id: str) -> TaskResponse:
         r = await self._client.get(
             f"{self.url}/api/v1/task/{task_id}",
@@ -97,7 +121,22 @@ class ChunkrAsync(ChunkrBase):
         )
         r.raise_for_status()
         return TaskResponse(**r.json()).with_client(self)
+    
+    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()
+    
+    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()
 
+    
     async def __aenter__(self):
         return self
 

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai/api/config.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, model_validator, ConfigDict
 from enum import Enum
 from typing import Optional, List, Dict
 
@@ -10,30 +10,30 @@ class CroppingStrategy(str, Enum):
     ALL = "All" 
     AUTO = "Auto"
 
-class LlmConfig(BaseModel):
-    model: str
-    prompt: str
-    temperature: float = 0.0
-
 class GenerationConfig(BaseModel):
     html: Optional[GenerationStrategy] = None
-    llm: Optional[LlmConfig] = None
+    llm: Optional[str] = None
     markdown: Optional[GenerationStrategy] = None
     crop_image: Optional[CroppingStrategy] = None
 
 class SegmentProcessing(BaseModel):
-    title: Optional[GenerationConfig] = None
-    section_header: Optional[GenerationConfig] = None
-    text: Optional[GenerationConfig] = None
-    list_item: Optional[GenerationConfig] = None
-    table: Optional[GenerationConfig] = None
-    picture: Optional[GenerationConfig] = None
-    caption: Optional[GenerationConfig] = None
-    formula: Optional[GenerationConfig] = None
-    footnote: Optional[GenerationConfig] = None
-    page_header: Optional[GenerationConfig] = None
-    page_footer: Optional[GenerationConfig] = None
-    page: Optional[GenerationConfig] = None
+    model_config = ConfigDict(
+        populate_by_name=True,
+        alias_generator=str.title
+    )
+    
+    title: Optional[GenerationConfig] = Field(default=None, alias="Title")
+    section_header: Optional[GenerationConfig] = Field(default=None, alias="SectionHeader")
+    text: Optional[GenerationConfig] = Field(default=None, alias="Text")
+    list_item: Optional[GenerationConfig] = Field(default=None, alias="ListItem")
+    table: Optional[GenerationConfig] = Field(default=None, alias="Table")
+    picture: Optional[GenerationConfig] = Field(default=None, alias="Picture")
+    caption: Optional[GenerationConfig] = Field(default=None, alias="Caption")
+    formula: Optional[GenerationConfig] = Field(default=None, alias="Formula")
+    footnote: Optional[GenerationConfig] = Field(default=None, alias="Footnote")
+    page_header: Optional[GenerationConfig] = Field(default=None, alias="PageHeader")
+    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
@@ -86,9 +86,9 @@ class Segment(BaseModel):
     bbox: BoundingBox
     content: str
     page_height: float
-    html: Optional[str]
-    image: Optional[str]
-    markdown: Optional[str]
+    html: Optional[str] = None
+    image: Optional[str] = None
+    markdown: Optional[str] = None
     ocr: List[OCRResult]
     page_number: int
     page_width: float
@@ -104,8 +104,8 @@ class ExtractedJson(BaseModel):
     data: Dict
 
 class OutputResponse(BaseModel):
-    chunks: List[Chunk] = []
-    extracted_json: Optional[ExtractedJson]
+    chunks: List[Chunk]
+    extracted_json: Optional[ExtractedJson] = Field(default=None)
 
 class Model(str, Enum):
     FAST = "Fast"

chunkr_ai-0.0.8/src/chunkr_ai/api/misc.py

@@ -0,0 +1,106 @@
+import io
+import json
+from pathlib import Path
+from PIL import Image
+import requests
+from typing import Union, Tuple, BinaryIO, Optional
+from .config import Configuration
+
+
+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
+    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(
+    file: Optional[Union[str, Path, BinaryIO, Image.Image]] = None,
+    config: Optional[Configuration] = None
+) -> dict:
+    """Prepare files and data dictionaries for upload.
+    
+    Args:
+        file: The file to upload
+        config: Optional configuration settings
+        
+    Returns:
+        dict: (files dict) ready for upload
+    """
+    files = {}
+    if file:
+        filename, file_obj = prepare_file(file)
+        files = {"file": (filename, file_obj)}
+
+    if config:
+        config_dict = config.model_dump(mode="json", exclude_none=True)
+        for key, value in config_dict.items():
+            files[key] = (None, json.dumps(value), 'application/json')
+                
+    return files

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai/api/task.py

@@ -1,5 +1,6 @@
 from .protocol import ChunkrClientProtocol
 from .config import Configuration, OutputResponse
+from .misc import prepare_upload_data
 import asyncio
 from datetime import datetime
 from enum import Enum
@@ -12,22 +13,23 @@ class Status(str, Enum):
     PROCESSING = "Processing"
     SUCCEEDED = "Succeeded"
     FAILED = "Failed"
+    CANCELLED = "Cancelled"
     
 class TaskResponse(BaseModel):
     configuration: Configuration
     created_at: datetime
-    expires_at: Optional[datetime]
-    file_name: Optional[str]
-    finished_at: Optional[datetime]
-    input_file_url: Optional[str]
+    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]
-    page_count: Optional[int]
-    pdf_url: Optional[str]
-    started_at: Optional[datetime]
+    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]
+    task_url: Optional[str] = None
     _client: Optional[Union[ChunkrClientProtocol]] = PrivateAttr(default=None)
 
     def with_client(self, client: Union[ChunkrClientProtocol]) -> 'TaskResponse':
@@ -79,7 +81,8 @@ class TaskResponse(BaseModel):
         """Poll the task for completion."""
         while True:
             response = self._poll_request_sync()
-            self.__dict__.update(response)
+            updated_task = TaskResponse(**response).with_client(self._client)
+            self.__dict__.update(updated_task.__dict__)
             
             if result := self._check_status():
                 return result
@@ -90,7 +93,8 @@ class TaskResponse(BaseModel):
         """Poll the task for completion asynchronously."""
         while True:
             response = await self._poll_request_async()
-            self.__dict__.update(response)
+            updated_task = TaskResponse(**response).with_client(self._client)
+            self.__dict__.update(updated_task.__dict__)
             
             if result := self._check_status():
                 return result
@@ -108,6 +112,56 @@ class TaskResponse(BaseModel):
                 if content:
                     parts.append(content)
         return "\n".join(parts)
+    
+    def update(self, config: Configuration) -> 'TaskResponse':
+        files = prepare_upload_data(None, config)
+        r = self._client._session.patch(
+            f"{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)
+    
+    def cancel(self):
+        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):
+        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}",
+            headers=self._client._headers()
+        )
+        r.raise_for_status()
 
     def html(self) -> str:
         """Get full HTML for the task"""

chunkr_ai-0.0.8/src/chunkr_ai/main.py

@@ -0,0 +1,12 @@
+from chunkr_ai.api.chunkr import Chunkr
+from chunkr_ai.models import Configuration
+from chunkr_ai.api.config import SegmentationStrategy, ChunkProcessing
+
+if __name__ == "__main__":
+    chunkr = Chunkr()
+    task = chunkr.update_task("556b4fe5-e3f7-48dc-9f56-0fb7fbacdb87", Configuration(
+        chunk_processing=ChunkProcessing(
+            target_length=1000
+        )
+    ))
+    print(task)

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai/models.py

@@ -8,7 +8,6 @@ from .api.config import (
     GenerationStrategy,
     GenerationConfig,
     JsonSchema,
-    LlmConfig,
     Model,
     OCRResult,
     OcrStrategy,

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8/src/chunkr_ai.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: chunkr-ai
-Version: 0.0.6
+Version: 0.0.8
 Summary: Python client for Chunkr: open source document intelligence
 Author-email: Ishaan Kapoor <ishaan@lumina.sh>
 Project-URL: Homepage, https://chunkr.ai

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/src/chunkr_ai.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ src/chunkr_ai/api/base.py
 src/chunkr_ai/api/chunkr.py
 src/chunkr_ai/api/chunkr_async.py
 src/chunkr_ai/api/config.py
+src/chunkr_ai/api/misc.py
 src/chunkr_ai/api/protocol.py
 src/chunkr_ai/api/task.py
 tests/test_chunkr.py

{chunkr_ai-0.0.6 → chunkr_ai-0.0.8}/tests/test_chunkr.py

@@ -19,7 +19,7 @@ from chunkr_ai.models import (
 
 @pytest.fixture(params=[
     pytest.param(("sync", Chunkr()), id="sync"),
-    pytest.param(("async", ChunkrAsync()), id="async")
+    # pytest.param(("async", ChunkrAsync()), id="async")
 ])
 def chunkr_client(request):
     return request.param
@@ -209,4 +209,121 @@ async def test_json_schema(chunkr_client, sample_path):
     assert response.status == "Succeeded"
     assert response.output is not None
     
-    
+@pytest.mark.asyncio
+async def test_delete_task(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    response = await client.upload(sample_path) if client_type == "async" else client.upload(sample_path)
+    assert isinstance(response, TaskResponse)
+    assert response.task_id is not None
+    assert response.status == "Succeeded"
+    assert response.output is not None
+    
+    if client_type == "async":
+        await client.delete_task(response.task_id)
+        with pytest.raises(Exception):  
+            await client.get_task(response.task_id)
+    else:
+        client.delete_task(response.task_id)
+        with pytest.raises(Exception):  
+            client.get_task(response.task_id)
+            
+@pytest.mark.asyncio
+async def test_delete_task_direct(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    task = await client.upload(sample_path) if client_type == "async" else client.upload(sample_path)
+    assert isinstance(task, TaskResponse)
+    assert task.task_id is not None
+    assert task.status == "Succeeded"
+    assert task.output is not None
+    
+    if client_type == "async":
+        await client.delete_task(task.task_id)
+        with pytest.raises(Exception):  
+            await client.get_task(task.task_id)
+    else:
+        client.delete_task(task.task_id)
+        with pytest.raises(Exception):  
+            client.get_task(task.task_id)
+    
+@pytest.mark.asyncio
+async def test_cancel_task(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    response = await client.create_task(sample_path) if client_type == "async" else client.create_task(sample_path)
+    assert isinstance(response, TaskResponse)
+    assert response.task_id is not None
+    assert response.status == "Starting"
+    
+    if client_type == "async":
+        await client.cancel_task(response.task_id)
+        assert (await client.get_task(response.task_id)).status == "Cancelled"
+        await response.poll_async()
+    else:
+        client.cancel_task(response.task_id)
+        assert client.get_task(response.task_id).status == "Cancelled"
+        response.poll()
+        
+    assert response.output is None
+    
+@pytest.mark.asyncio
+async def test_cancel_task_direct(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    task = await client.create_task(sample_path) if client_type == "async" else client.create_task(sample_path)
+    assert isinstance(task, TaskResponse)
+    assert task.task_id is not None
+    assert task.status == "Starting"
+    
+    if client_type == "async":
+        await task.cancel_async()
+    else:
+        task.cancel()
+    
+    assert task.status == "Cancelled"
+    assert task.output is None
+    
+@pytest.mark.asyncio
+async def test_update_task(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    original_config = Configuration(
+        segmentation_strategy=SegmentationStrategy.LAYOUT_ANALYSIS,
+    )
+    new_config = Configuration(
+        segmentation_strategy=SegmentationStrategy.PAGE,
+    )
+    response = await client.upload(sample_path, original_config) if client_type == "async" else client.upload(sample_path, original_config)
+    assert isinstance(response, TaskResponse)
+    assert response.task_id is not None
+    assert response.status == "Succeeded"
+    assert response.output is not None
+    
+    if client_type == "async":
+        task = await client.update(response.task_id, new_config)
+    else:
+        task = client.update(response.task_id, new_config)
+
+    assert task.status == "Succeeded"
+    assert task.output is not None
+    assert task.configuration.segmentation_strategy == SegmentationStrategy.PAGE
+    
+@pytest.mark.asyncio
+async def test_update_task_direct(chunkr_client, sample_path):
+    client_type, client = chunkr_client
+    original_config = Configuration(
+        segmentation_strategy=SegmentationStrategy.LAYOUT_ANALYSIS,
+    )
+    new_config = Configuration(
+        segmentation_strategy=SegmentationStrategy.PAGE,
+    )
+    task = await client.upload(sample_path, original_config) if client_type == "async" else client.upload(sample_path, original_config)
+    assert isinstance(task, TaskResponse)
+    assert task.task_id is not None
+    assert task.status == "Succeeded"
+    assert task.output is not None
+    
+    if client_type == "async":
+        await task.update_async(new_config)
+    else:
+        task.update(new_config)
+
+    assert task.status == "Succeeded"
+    assert task.output is not None
+    assert task.configuration.segmentation_strategy == SegmentationStrategy.PAGE

chunkr_ai-0.0.6/src/chunkr_ai/api/base.py

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