tests, file improvements, download and dir helpers

okaris · okaris · commit 04579646cbe9 · 2025-05-06T08:38:27.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,3 +35,9 @@ addopts = "-v"
 [tool.flake8]
 max-line-length = 100
 exclude = [".git", "__pycache__", "build", "dist"]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+]
diff --git a/src/inferencesh/sdk.py b/src/inferencesh/sdk.py
@@ -1,17 +1,20 @@
 from typing import Optional, Union
-from pydantic import BaseModel, ConfigDict, PrivateAttr, model_validator
+from pydantic import BaseModel, ConfigDict, PrivateAttr, model_validator, Field, field_validator
 import mimetypes
 import os
 import urllib.request
 import urllib.parse
 import tempfile
-from pydantic import Field
 from typing import Any, Dict, List
 
 import inspect
 import ast
 import textwrap
 from collections import OrderedDict
+from enum import Enum
+import shutil
+from pathlib import Path
+import hashlib
 
 
 # inspired by https://github.com/pydantic/pydantic/issues/7580
@@ -102,39 +105,35 @@ async def unload(self):
 
 class File(BaseModel):
     """A class representing a file in the inference.sh ecosystem."""
-    uri: Optional[str] = None  # Original location (URL or file path)
+    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
     size: Optional[int] = None  # File size in bytes
     filename: Optional[str] = None  # Original filename if available
     _tmp_path: Optional[str] = PrivateAttr(default=None)  # Internal storage for temporary file path
     
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-        populate_by_name=True
-    )
-
-    @classmethod
-    def __get_validators__(cls):
-        # First yield the default validators
-        yield cls.validate
-
+    def __init__(self, initializer=None, **data):
+        if initializer is not None:
+            if isinstance(initializer, str):
+                data['uri'] = initializer
+            elif isinstance(initializer, File):
+                data = initializer.model_dump()
+            else:
+                raise ValueError(f'Invalid input for File: {initializer}')
+        super().__init__(**data)
+
+    @model_validator(mode='before')
     @classmethod
-    def validate(cls, value):
-        """Convert string values to File objects."""
-        if isinstance(value, str):
-            # If it's a string, treat it as a uri
-            return cls(uri=value)
-        elif isinstance(value, cls):
-            # If it's already a File instance, return it as is
-            return value
-        elif isinstance(value, dict):
-            # If it's a dict, use normal validation
-            return cls(**value)
-        raise ValueError(f'Invalid input for File: {value}')
-
+    def convert_str_to_file(cls, values):
+        print(f"check_uri_or_path input: {values}")
+        if isinstance(values, str):  # Only accept strings
+            return {"uri": values}
+        elif isinstance(values, dict):
+            return values
+        raise ValueError(f'Invalid input for File: {values}')
+    
     @model_validator(mode='after')
-    def check_uri_or_path(self) -> 'File':
+    def validate_required_fields(self) -> 'File':
         """Validate that either uri or path is provided."""
         if not self.uri and not self.path:
             raise ValueError("Either 'uri' or 'path' must be provided")
@@ -147,7 +146,10 @@ def model_post_init(self, _: Any) -> None:
             self.path = os.path.abspath(self.uri)
         elif self.uri:
             self.path = self.uri
-        self._populate_metadata()
+        if self.path:
+            self._populate_metadata()
+        else:
+            raise ValueError("Either 'uri' or 'path' must be provided")
     
     def _is_url(self, path: str) -> bool:
         """Check if the path is a URL."""
@@ -234,13 +236,24 @@ def exists(self) -> bool:
     
     def refresh_metadata(self) -> None:
         """Refresh all metadata from the file."""
-        self._populate_metadata() 
+        if os.path.exists(self.path):
+            self.content_type = self._guess_content_type()
+            self.size = self._get_file_size()  # Always update size
+            self.filename = self._get_filename()
 
 
+class ContextMessageRole(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant" 
+    SYSTEM = "system"
+
+class Message(BaseModel):
+    role: ContextMessageRole
+    content: str
+
 class ContextMessage(BaseModel):
-    role: str = Field(
+    role: ContextMessageRole = Field(
         description="The role of the message",
-        enum=["user", "assistant", "system"]
     )
     text: str = Field(
         description="The text content of the message"
@@ -300,4 +313,51 @@ class LLMInputWithImage(LLMInput):
     image: Optional[File] = Field(
         description="The image to use for the model",
         default=None
-    )
+    )
+
+class DownloadDir(str, Enum):
+    """Standard download directories used by the SDK."""
+    DATA = "./data"   # Persistent storage/cache directory
+    TEMP = "./tmp"    # Temporary storage directory
+    CACHE = "./cache" # Cache directory
+
+def download(url: str, directory: Union[str, Path, DownloadDir]) -> str:
+    """Download a file to the specified directory and return its path.
+    
+    Args:
+        url: The URL to download from
+        directory: The directory to save the file to. Can be a string path, 
+                  Path object, or DownloadDir enum value.
+        
+    Returns:
+        str: The path to the downloaded file
+    """
+    # Convert directory to Path
+    dir_path = Path(directory)
+    dir_path.mkdir(exist_ok=True)
+    
+    # Create hash directory from URL
+    url_hash = hashlib.sha256(url.encode()).hexdigest()[:12]
+    hash_dir = dir_path / url_hash
+    hash_dir.mkdir(exist_ok=True)
+    
+    # Keep original filename
+    filename = os.path.basename(urllib.parse.urlparse(url).path)
+    if not filename:
+        filename = 'download'
+        
+    output_path = hash_dir / filename
+    
+    # If file exists in directory and it's not a temp directory, return it
+    if output_path.exists() and directory != DownloadDir.TEMP:
+        return str(output_path)
+    
+    # Download the file
+    file = File(url)
+    if file.path:
+        shutil.copy2(file.path, output_path)
+        # Prevent the File instance from deleting its temporary file
+        file._tmp_path = None
+        return str(output_path)
+    
+    raise RuntimeError(f"Failed to download {url}")
diff --git a/tests/test_sdk.py b/tests/test_sdk.py
@@ -1,6 +1,8 @@
 import os
 import pytest
+import tempfile
 from inferencesh import BaseApp, BaseAppInput, BaseAppOutput, File
+import urllib.parse
 
 def test_file_creation():
     # Create a temporary file
@@ -22,10 +24,139 @@ class TestInput(BaseAppInput):
     class TestOutput(BaseAppOutput):
         result: str
 
-    class TestApp(BaseApp):
-        async def run(self, app_input: TestInput) -> TestOutput:
-            return TestOutput(result=f"Processed: {app_input.text}")
-
-    app = TestApp()
+    # Use BaseApp directly, don't subclass with implementation
+    app = BaseApp()
+    import asyncio
     with pytest.raises(NotImplementedError):
-        app.run(TestInput(text="test")) 
+        asyncio.run(app.run(TestInput(text="test")))
+
+def test_file_from_local_path():
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(suffix='.txt', delete=False) as f:
+        f.write(b"test content")
+        path = f.name
+    
+    try:
+        # Test creating File from path
+        file = File(uri=path)
+        assert file.exists()
+        assert file.size == len("test content")
+        assert file.content_type == "text/plain"
+        assert file.filename == os.path.basename(path)
+        assert file.path == os.path.abspath(path)
+        assert file._tmp_path is None  # Should not create temp file for local paths
+    finally:
+        os.unlink(path)
+
+def test_file_from_relative_path():
+    # Create a file in current directory
+    with open("test_relative.txt", "w") as f:
+        f.write("relative test")
+    
+    try:
+        file = File(uri="test_relative.txt")
+        assert file.exists()
+        assert os.path.isabs(file.path)
+        assert file.filename == "test_relative.txt"
+    finally:
+        os.unlink("test_relative.txt")
+
+def test_file_validation():
+    # Test empty initialization
+    with pytest.raises(ValueError, match="Either 'uri' or 'path' must be provided"):
+        File()
+    
+    # Test invalid input type
+    with pytest.raises(ValueError, match="Invalid input for File"):
+        File(123)
+    
+    # Test string input (should work)
+    with tempfile.NamedTemporaryFile(suffix='.txt', delete=False) as f:
+        f.write(b"test content")
+        path = f.name
+        
+    try:
+        file = File(path)
+        assert isinstance(file, File)
+        assert file.uri == path
+        assert file.exists()
+    finally:
+        os.unlink(path)
+
+def test_file_from_url(monkeypatch):
+    # Mock URL download
+    def mock_urlopen(request):
+        class MockResponse:
+            def __enter__(self):
+                return self
+            
+            def __exit__(self, *args):
+                pass
+            
+            def read(self):
+                return b"mocked content"
+        
+        return MockResponse()
+    
+    monkeypatch.setattr(urllib.request, 'urlopen', mock_urlopen)
+    
+    url = "https://example.com/test.pdf"
+    file = File(uri=url)
+    
+    try:
+        assert file._is_url(url)
+        assert file.exists()
+        assert file._tmp_path is not None
+        assert file._tmp_path.endswith('.pdf')  # Just check the extension
+        assert file.content_type == "application/pdf"
+    finally:
+        # Cleanup should happen in __del__ but let's be explicit for testing
+        if file._tmp_path and os.path.exists(file._tmp_path):
+            os.unlink(file._tmp_path)
+
+def test_file_metadata_refresh():
+    with tempfile.NamedTemporaryFile(suffix='.json', delete=False) as f:
+        initial_content = b'{"test": "data"}'
+        f.write(initial_content)
+        path = f.name
+    
+    try:
+        file = File(uri=path)
+        initial_size = file.size
+        
+        # Modify file with significantly more data
+        with open(path, 'ab') as f:  # Open in append binary mode
+            additional_data = b'\n{"more": "data"}\n' * 10  # Add multiple lines of data
+            f.write(additional_data)
+        
+        # Refresh metadata
+        file.refresh_metadata()
+        assert file.size > initial_size, f"New size {file.size} should be larger than initial size {initial_size}"
+    finally:
+        os.unlink(path)
+
+def test_file_cleanup(monkeypatch):
+    # Mock URL download - same mock as test_file_from_url
+    def mock_urlopen(request):
+        class MockResponse:
+            def __enter__(self):
+                return self
+            
+            def __exit__(self, *args):
+                pass
+            
+            def read(self):
+                return b"mocked content"
+        
+        return MockResponse()
+    
+    monkeypatch.setattr(urllib.request, 'urlopen', mock_urlopen)
+    
+    url = "https://example.com/test.txt"
+    file = File(uri=url)
+    
+    if file._tmp_path:
+        tmp_path = file._tmp_path
+        assert os.path.exists(tmp_path)
+        del file
+        assert not os.path.exists(tmp_path) 
