inferencesh 0.2.23__py3-none-any.whl → 0.4.29__py3-none-any.whl

inferencesh/models/base.py

@@ -1,10 +1,23 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 from pydantic import BaseModel, ConfigDict
 import inspect
 import ast
 import textwrap
 from collections import OrderedDict
-
+from inferencesh.models.file import File
+from pydantic import Field
+
+class Metadata(BaseModel):
+    app_id: Optional[str] = None
+    app_version_id: Optional[str] = None
+    app_variant: Optional[str] = None
+    worker_id: Optional[str] = None
+    def update(self, other: Dict[str, Any] | BaseModel) -> None:
+        update_dict = other.model_dump() if isinstance(other, BaseModel) else other
+        for key, value in update_dict.items():
+            setattr(self, key, value)
+    class Config:
+        extra = "allow"
 
 class OrderedSchemaModel(BaseModel):
     """A base model that ensures the JSON schema properties and required fields are in the order of field definition."""
@@ -91,4 +104,69 @@ class BaseApp(BaseModel):
         raise NotImplementedError("run method must be implemented")
 
     async def unload(self):
-        pass 
+        pass 
+    
+    
+# Mixins
+
+class OptionalImageFieldMixin(BaseModel):
+    image: Optional[File] = Field(
+        description="the image to use for the model",
+        default=None,
+        contentMediaType="image/*",
+    )
+
+class RequiredImageFieldMixin(BaseModel):
+    image: File = Field(
+        description="the image to use for the model",
+        contentMediaType="image/*",
+    )
+    
+class OptionalVideoFieldMixin(BaseModel):
+    video: Optional[File] = Field(
+        description="the video to use for the model",
+        default=None,
+        contentMediaType="video/*",
+    )
+    
+class RequiredVideoFieldMixin(BaseModel):
+    video: File = Field(
+        description="the video to use for the model",
+        contentMediaType="video/*",
+    )
+    
+class OptionalAudioFieldMixin(BaseModel):
+    audio: Optional[File] = Field(
+        description="the audio to use for the model",
+        default=None,
+        contentMediaType="audio/*",
+    )
+    
+class RequiredAudioFieldMixin(BaseModel):
+    audio: File = Field(
+        description="the audio to use for the model",
+        contentMediaType="audio/*",
+    )
+    
+class OptionalTextFieldMixin(BaseModel):
+    text: Optional[str] = Field(
+        description="the text to use for the model",
+        default=None,
+    )
+    
+class RequiredTextFieldMixin(BaseModel):
+    text: str = Field(
+        description="the text to use for the model",
+    )
+    
+class OptionalFileFieldMixin(BaseModel):
+    file: Optional[File] = Field(
+        description="the file to use for the model",
+        default=None,
+    )
+    
+class RequiredFileFieldMixin(BaseModel):
+    file: Optional[File] = Field(
+        description="the file to use for the model",
+        default=None,
+    )

inferencesh/models/file.py

@@ -1,15 +1,48 @@
 from typing import Optional, Union, Any
-from pydantic import BaseModel, Field, PrivateAttr, model_validator
+from pydantic import BaseModel, Field, PrivateAttr, model_validator, GetJsonSchemaHandler
+from pydantic_core import CoreSchema
 import mimetypes
 import os
 import urllib.request
 import urllib.parse
-import tempfile
+import hashlib
+from pathlib import Path
 from tqdm import tqdm
 
-
 class File(BaseModel):
     """A class representing a file in the inference.sh ecosystem."""
+    
+    @classmethod
+    def get_cache_dir(cls) -> Path:
+        """Get the cache directory path based on environment variables or default location."""
+        if cache_dir := os.environ.get("FILE_CACHE_DIR"):
+            path = Path(cache_dir)
+        else:
+            path = Path.home() / ".cache" / "inferencesh" / "files"
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+    
+    def _get_cache_path(self, url: str) -> Path:
+        """Get the cache path for a URL using a hash-based directory structure."""
+        # Parse URL components
+        parsed_url = urllib.parse.urlparse(url)
+        
+        # Create hash from URL path and query parameters for uniqueness
+        url_components = parsed_url.netloc + parsed_url.path
+        if parsed_url.query:
+            url_components += '?' + parsed_url.query
+        url_hash = hashlib.sha256(url_components.encode()).hexdigest()[:12]
+        
+        # Get filename from URL or use default
+        filename = os.path.basename(parsed_url.path)
+        if not filename:
+            filename = 'download'
+            
+        # Create hash directory in cache
+        cache_dir = self.get_cache_dir() / url_hash
+        cache_dir.mkdir(exist_ok=True)
+        
+        return cache_dir / filename
     uri: Optional[str] = Field(default=None)  # Original location (URL or file path)
     path: Optional[str] = None  # Resolved local file path
     content_type: Optional[str] = None  # MIME type of the file
@@ -74,14 +107,21 @@ class File(BaseModel):
         return parsed.scheme in ('http', 'https')
 
     def _download_url(self) -> None:
-        """Download the URL to a temporary file and update the path."""
+        """Download the URL to the cache directory and update the path."""
         original_url = self.uri
-        tmp_file = None
+        cache_path = self._get_cache_path(original_url)
+        
+        # If file exists in cache, use it
+        if cache_path.exists():
+            print(f"Using cached file: {cache_path}")
+            self.path = str(cache_path)
+            return
+            
+        print(f"Downloading URL: {original_url} to {cache_path}")
         try:
-            # Create a temporary file with a suffix based on the URL path
-            suffix = os.path.splitext(urllib.parse.urlparse(original_url).path)[1]
-            tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
-            self._tmp_path = tmp_file.name
+            # Download to a temporary filename in the final directory
+            tmp_path = str(cache_path) + '.tmp'
+            self._tmp_path = tmp_path
             
             # Set up request with user agent
             headers = {
@@ -97,26 +137,53 @@ class File(BaseModel):
             print(f"Downloading URL: {original_url} to {self._tmp_path}")
             try:
                 with urllib.request.urlopen(req) as response:
-                    total_size = int(response.headers.get('content-length', 0))
+                    # Safely retrieve content-length if available
+                    total_size = 0
+                    try:
+                        if hasattr(response, 'headers') and response.headers is not None:
+                            # urllib may expose headers as an email.message.Message
+                            cl = response.headers.get('content-length')
+                            total_size = int(cl) if cl is not None else 0
+                        elif hasattr(response, 'getheader'):
+                            cl = response.getheader('content-length')
+                            total_size = int(cl) if cl is not None else 0
+                    except Exception:
+                        total_size = 0
+
                     block_size = 1024  # 1 Kibibyte
                     
                     with tqdm(total=total_size, unit='iB', unit_scale=True) as pbar:
                         with open(self._tmp_path, 'wb') as out_file:
                             while True:
-                                buffer = response.read(block_size)
+                                non_chunking = False
+                                try:
+                                    buffer = response.read(block_size)
+                                except TypeError:
+                                    # Some mocks (or minimal implementations) expose read() without size
+                                    buffer = response.read()
+                                    non_chunking = True
                                 if not buffer:
                                     break
                                 out_file.write(buffer)
-                                pbar.update(len(buffer))
+                                try:
+                                    pbar.update(len(buffer))
+                                except Exception:
+                                    pass
+                                if non_chunking:
+                                    # If we read the whole body at once, exit loop
+                                    break
                             
-                self.path = self._tmp_path
+                # Rename the temporary file to the final name
+                os.rename(self._tmp_path, cache_path)
+                self._tmp_path = None  # Prevent deletion in __del__
+                self.path = str(cache_path)
             except (urllib.error.URLError, urllib.error.HTTPError) as e:
                 raise RuntimeError(f"Failed to download URL {original_url}: {str(e)}")
             except IOError as e:
                 raise RuntimeError(f"Failed to write downloaded file to {self._tmp_path}: {str(e)}")
         except Exception as e:
             # Clean up temp file if something went wrong
-            if tmp_file is not None and hasattr(self, '_tmp_path'):
+            if hasattr(self, '_tmp_path') and self._tmp_path:
                 try:
                     os.unlink(self._tmp_path)
                 except (OSError, IOError):
@@ -169,14 +236,46 @@ class File(BaseModel):
             self.size = self._get_file_size()  # Always update size
             self.filename = self._get_filename()
 
+    # @classmethod
+    # def __get_pydantic_core_schema__(
+    #     cls, source: Type[Any], handler: GetCoreSchemaHandler
+    # ) -> CoreSchema:
+    #     """Generates a Pydantic Core schema for validation of this File class"""
+    #     # Get the default schema for our class
+    #     schema = handler(source)
+        
+    #     # Create a proper serialization schema that includes the type
+    #     serialization = core_schema.plain_serializer_function_ser_schema(
+    #         lambda x: x.uri if x.uri else x.path,
+    #         return_schema=core_schema.str_schema(),
+    #         when_used="json",
+    #     )
+        
+    #     return core_schema.json_or_python_schema(
+    #         json_schema=core_schema.union_schema([
+    #             core_schema.str_schema(),  # Accept string input
+    #             schema,  # Accept full object input
+    #         ]),
+    #         python_schema=schema,
+    #         serialization=serialization,
+    #     )
+
     @classmethod
-    def model_json_schema(cls, **kwargs):
-        schema = super().model_json_schema(**kwargs)
-        schema["$id"] = "/schemas/File"
-        # Create a schema that accepts either a string or the full object
+    def __get_pydantic_json_schema__(
+        cls, schema: CoreSchema, handler: GetJsonSchemaHandler
+    ) -> dict[str, Any]:
+        """Generate a simple JSON schema that accepts either a string or an object"""
+        json_schema = handler(schema)
+        if "$ref" in json_schema:
+            # If we got a ref, resolve it to the actual schema
+            json_schema = handler.resolve_ref_schema(json_schema)
+        
+        # Add string as an alternative without recursion
         return {
+            "$id": "/schemas/File",
             "oneOf": [
-                {"type": "string"},  # Accept string input
-                schema  # Accept full object input
+                {k: v for k, v in json_schema.items() if k != "$ref"},  # Remove any $ref to prevent recursion
+                {"type": "string"}
             ]
-        } 
+        }
+