Support presigned URL for file upload (#102)

Author: spillai

vlmrun/client/files.py

@@ -2,14 +2,16 @@
 
 from __future__ import annotations
 
+import requests
+import time
 import hashlib
 from pathlib import Path
-from typing import Optional, Union, List
+from typing import Optional, Union, List, Literal
 
 from loguru import logger
 from vlmrun.client.base_requestor import APIRequestor
 from vlmrun.types.abstract import VLMRunProtocol
-from vlmrun.client.types import FileResponse
+from vlmrun.client.types import FileResponse, PresignedUrlResponse
 
 
 class Files:
@@ -89,40 +91,116 @@ def upload(
         self,
         file: Union[Path, str],
         purpose: str = "assistants",
+        method: Literal["auto", "direct", "presigned-url"] = "auto",
         timeout: int = 3 * 60,
+        expiration: int = 24 * 60 * 60,
+        force: bool = False,
     ) -> FileResponse:
         """Upload a file.
 
         Args:
             file: Path to file to upload
             purpose: Purpose of file (default: fine-tune)
             timeout: Timeout for upload (default: 3 minutes)
+            method: Method to use for upload (default: "auto")
+            expiration: Expiration time for presigned URL (default: 24 hours)
+            force: Force upload even if file already exists (default: False)
 
         Returns:
             FileResponse: Uploaded file object
         """
         if isinstance(file, str):
-            file = Path(file)
+            try:
+                file = Path(file)
+            except Exception as exc:
+                raise ValueError(f"Invalid file path: {file}") from exc
+        elif not isinstance(file, Path):
+            raise ValueError(f"Invalid file path: {file}")
+
+        # Check if the file exists
+        if not file.exists():
+            raise FileNotFoundError(f"File does not exist: {file}")
 
         # Check if the file already exists in the database
         cached_response: Optional[FileResponse] = self.get_cached_file(file)
-        if cached_response:
+        if cached_response and not force:
             return cached_response
 
-        # Upload the file
-        with open(file, "rb") as f:
-            files = {"file": (file.name, f)}
+        # If method is "auto", check if the file is too large to upload directly
+        if method == "auto":
+            if file.stat().st_size > 32 * 1024 * 1024:  # 32MB
+                method = "presigned-url"
+            else:
+                method = "direct"
+            logger.debug(f"Upload method [file={file}, method={method}]")
+
+        # If method is "presigned-url", generate a presigned URL for the file
+        if method == "presigned-url":
+            # Generate a presigned URL for the file
             response, status_code, headers = self._requestor.request(
                 method="POST",
-                url="files",
-                params={"purpose": purpose},
-                files=files,
-                timeout=timeout,
+                url="files/presigned-url",
+                data={
+                    "filename": file.name,
+                    "purpose": purpose,
+                    "expiration": expiration,
+                },
             )
-
             if not isinstance(response, dict):
                 raise TypeError("Expected dict response")
-            return FileResponse(**response)
+            response = PresignedUrlResponse(**response)
+
+            # PUT the file to the presigned URL
+            start_t = time.time()
+            logger.debug(
+                f"Uploading file to presigned URL [file={file}, id={response.id}, url={response.url}]"
+            )
+            with file.open("rb") as f:
+                put_response = requests.put(
+                    response.url,
+                    headers={"Content-Type": response.content_type},
+                    data=f,
+                )
+                status_code = put_response.status_code
+            end_t = time.time()
+            logger.debug(
+                f"Uploaded file to presigned URL [file={file}, url={response.url}, time={end_t - start_t:.1f}s]"
+            )
+            if status_code == 200:
+                # Verify the file upload
+                verify_response, status_code, headers = self._requestor.request(
+                    method="GET",
+                    url=f"files/verify-upload/{response.id}",
+                )
+                if status_code == 200:
+                    return FileResponse(**verify_response)
+                else:
+                    raise Exception(f"Failed to verify file upload: {verify_response}")
+            else:
+                raise Exception(f"Failed to upload file to presigned URL: {response}")
+
+        # If method is "direct", upload the file directly
+        elif method == "direct":
+            logger.debug(f"Uploading file directly [file={file}]")
+            # Upload the file
+            with open(file, "rb") as f:
+                files = {"file": (file.name, f)}
+                response, status_code, headers = self._requestor.request(
+                    method="POST",
+                    url="files",
+                    params={"purpose": purpose},
+                    files=files,
+                    timeout=timeout,
+                )
+            if status_code == 201:
+                if not isinstance(response, dict):
+                    raise TypeError("Expected dict response")
+                return FileResponse(**response)
+            else:
+                raise Exception(f"Failed to upload file directly: {response}")
+
+        else:
+            raise ValueError(f"Invalid upload method: {method}")
 
     def get(self, file_id: str) -> FileResponse:
         """Get file metadata.

vlmrun/client/hub.py

@@ -96,7 +96,9 @@ def list_domains(self) -> List[HubDomainInfo]:
         except Exception as e:
             raise APIError(f"Failed to list domains: {str(e)}")
 
-    def get_schema(self, domain: str, gql_stmt: Optional[str] = None) -> HubSchemaResponse:
+    def get_schema(
+        self, domain: str, gql_stmt: Optional[str] = None
+    ) -> HubSchemaResponse:
         """Get the JSON schema for a given domain.
 
         Args:

vlmrun/client/types.py

@@ -35,6 +35,16 @@ class FileResponse(BaseModel):
     object: str = "file"
 
 
+class PresignedUrlResponse(BaseModel):
+    id: Optional[str]
+    url: Optional[str]
+    filename: Optional[str]
+    expiration: Optional[int]
+    method: Optional[str]
+    content_type: Optional[str]
+    created_at: datetime
+
+
 class CreditUsage(BaseModel):
     elements_processed: Optional[int] = None
     element_type: Optional[Literal["image", "page", "video", "audio"]] = None

vlmrun/common/image.py

@@ -12,10 +12,10 @@
 
 def _open_image_with_exif(path: Union[str, Path]) -> Image.Image:
     """Open an image and apply EXIF orientation if available.
-    
+
     Args:
         path: Path to the image file
-        
+
     Returns:
         PIL Image with EXIF orientation applied and converted to RGB
     """

vlmrun/version.py

@@ -1 +1 @@
-__version__ = "0.2.16"
+__version__ = "0.2.17"