Add support for files API endpoints

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

replicate/__init__.py

@@ -14,8 +14,9 @@
 async_paginate = _async_paginate
 
 collections = default_client.collections
-hardware = default_client.hardware
 deployments = default_client.deployments
+files = default_client.files
+hardware = default_client.hardware
 models = default_client.models
 predictions = default_client.predictions
 trainings = default_client.trainings

replicate/client.py

@@ -24,6 +24,7 @@
 from replicate.collection import Collections
 from replicate.deployment import Deployments
 from replicate.exceptions import ReplicateError
+from replicate.file import Files
 from replicate.hardware import HardwareNamespace as Hardware
 from replicate.model import Models
 from replicate.prediction import Predictions
@@ -117,6 +118,13 @@ def deployments(self) -> Deployments:
         """
         return Deployments(client=self)
 
+    @property
+    def files(self) -> Files:
+        """
+        Namespace for operations related to files.
+        """
+        return Files(client=self)
+
     @property
     def hardware(self) -> Hardware:
         """

replicate/file.py

@@ -0,0 +1,202 @@
+import base64
+import io
+import json
+import mimetypes
+import os
+import pathlib
+from typing import Any, BinaryIO, Dict, List, Optional, Tuple, TypedDict, Union
+
+import httpx
+from typing_extensions import NotRequired, Unpack
+
+from replicate.resource import Namespace, Resource
+
+
+class File(Resource):
+    """
+    A file uploaded to Replicate that can be used as an input to a model.
+    """
+
+    id: str
+    """The ID of the file."""
+
+    name: str
+    """The name of the file."""
+
+    content_type: str
+    """The content type of the file."""
+
+    size: int
+    """The size of the file in bytes."""
+
+    etag: str
+    """The ETag of the file."""
+
+    checksums: Dict[str, str]
+    """The checksums of the file."""
+
+    metadata: Dict[str, Any]
+    """The metadata of the file."""
+
+    created_at: str
+    """The time the file was created."""
+
+    expires_at: Optional[str]
+    """The time the file will expire."""
+
+    urls: Dict[str, str]
+    """The URLs of the file."""
+
+
+class Files(Namespace):
+    class CreateFileParams(TypedDict):
+        """Parameters for creating a file."""
+
+        filename: NotRequired[str]
+        """The name of the file."""
+
+        content_type: NotRequired[str]
+        """The content type of the file."""
+
+        metadata: NotRequired[Dict[str, Any]]
+        """The file metadata."""
+
+    def create(
+        self,
+        file: Union[str, pathlib.Path, BinaryIO, io.IOBase],
+        **params: Unpack["Files.CreateFileParams"],
+    ) -> File:
+        """
+        Upload a file that can be passed as an input when running a model.
+        """
+
+        if isinstance(file, (str, pathlib.Path)):
+            with open(file, "rb") as f:
+                return self.create(f, **params)
+        elif not isinstance(file, (io.IOBase, BinaryIO)):
+            raise ValueError(
+                "Unsupported file type. Must be a file path or file-like object."
+            )
+
+        resp = self._client._request(
+            "POST", "/v1/files", timeout=None, **_create_file_params(file, **params)
+        )
+
+        return _json_to_file(resp.json())
+
+    async def async_create(
+        self,
+        file: Union[str, pathlib.Path, BinaryIO, io.IOBase],
+        **params: Unpack["Files.CreateFileParams"],
+    ) -> File:
+        """Upload a file asynchronously that can be passed as an input when running a model."""
+
+        if isinstance(file, (str, pathlib.Path)):
+            with open(file, "rb") as f:
+                return self.create(f, **params)
+        elif not isinstance(file, (io.IOBase, BinaryIO)):
+            raise ValueError(
+                "Unsupported file type. Must be a file path or file-like object."
+            )
+
+        resp = await self._client._async_request(
+            "POST", "/v1/files", timeout=None, **_create_file_params(file, **params)
+        )
+
+        return _json_to_file(resp.json())
+
+    def get(self, file_id: str) -> File:
+        """Get an uploaded file by its ID."""
+
+        resp = self._client._request("GET", f"/v1/files/{file_id}")
+        return _json_to_file(resp.json())
+
+    async def async_get(self, file_id: str) -> File:
+        """Get an uploaded file by its ID asynchronously."""
+
+        resp = await self._client._async_request("GET", f"/v1/files/{file_id}")
+        return _json_to_file(resp.json())
+
+    def list(self) -> List[File]:
+        """List all uploaded files."""
+
+        resp = self._client._request("GET", "/v1/files")
+        return [_json_to_file(obj) for obj in resp.json().get("results", [])]
+
+    async def async_list(self) -> List[File]:
+        """List all uploaded files asynchronously."""
+
+        resp = await self._client._async_request("GET", "/v1/files")
+        return [_json_to_file(obj) for obj in resp.json().get("results", [])]
+
+    def delete(self, file_id: str) -> None:
+        """Delete an uploaded file by its ID."""
+
+        _ =self._client._request("DELETE", f"/v1/files/{file_id}")
+
+    async def async_delete(self, file_id: str) -> None:
+        """Delete an uploaded file by its ID asynchronously."""
+
+        _ = await self._client._async_request("DELETE", f"/v1/files/{file_id}")
+
+def _create_file_params(
+    file: Union[BinaryIO, io.IOBase],
+    **params: Unpack["Files.CreateFileParams"],
+) -> Dict[str, Any]:
+    file.seek(0)
+
+    if params is None:
+        params = {}
+
+    filename = params.get("filename", os.path.basename(getattr(file, "name", "file")))
+    content_type = (
+        params.get("content_type")
+        or mimetypes.guess_type(filename)[0]
+        or "application/octet-stream"
+    )
+    metadata = params.get("metadata")
+
+    data = {}
+    if metadata:
+        data["metadata"] = json.dumps(metadata)
+
+    return {
+        "files": {"content": (filename, file, content_type)},
+        "data": data,
+    }
+
+def _json_to_file(json: Dict[str, Any]) -> File:  # pylint: disable=redefined-outer-name
+    return File(**json)
+
+
+def upload_file(file: io.IOBase, output_file_prefix: Optional[str] = None) -> str:
+    """
+    Upload a file to the server.
+
+    Args:
+        file: A file handle to upload.
+        output_file_prefix: A string to prepend to the output file name.
+    Returns:
+        str: A URL to the uploaded file.
+    """
+    # Lifted straight from cog.files
+
+    file.seek(0)
+
+    if output_file_prefix is not None:
+        name = getattr(file, "name", "output")
+        url = output_file_prefix + os.path.basename(name)
+        resp = httpx.put(url, files={"file": file}, timeout=None)  # type: ignore
+        resp.raise_for_status()
+
+        return url
+
+    body = file.read()
+    # Ensure the file handle is in bytes
+    body = body.encode("utf-8") if isinstance(body, str) else body
+    encoded_body = base64.b64encode(body).decode("utf-8")
+    # Use getattr to avoid mypy complaints about io.IOBase having no attribute name
+    mime_type = (
+        mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
+    )
+    return f"data:{mime_type};base64,{encoded_body}"

replicate/files.py

@@ -19,7 +19,7 @@
 from typing_extensions import NotRequired, TypedDict, Unpack
 
 from replicate.exceptions import ModelError, ReplicateError
-from replicate.files import upload_file
+from replicate.file import upload_file
 from replicate.json import encode_json
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource

replicate/prediction.py

@@ -13,7 +13,7 @@
 
 from typing_extensions import NotRequired, Unpack
 
-from replicate.files import upload_file
+from replicate.file import upload_file
 from replicate.identifier import ModelVersionIdentifier
 from replicate.json import encode_json
 from replicate.model import Model

replicate/training.py

@@ -0,0 +1,58 @@
+import tempfile
+
+import pytest
+
+import replicate
+
+
+# @pytest.mark.vcr("files-operations.yaml")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_flag", [True, False])
+async def test_file_operations(async_flag):
+    # Create a sample file
+    with tempfile.NamedTemporaryFile(
+        mode="wb", delete=False, prefix="test_file", suffix=".txt"
+    ) as temp_file:
+        temp_file.write(b"Hello, Replicate!")
+
+        # Test create
+        if async_flag:
+            created_file = await replicate.files.async_create(temp_file.name)
+        else:
+            created_file = replicate.files.create(temp_file.name)
+
+    assert created_file.name.startswith("test_file")
+    assert created_file.name.endswith(".txt")
+    file_id = created_file.id
+
+    # Test get
+    if async_flag:
+        retrieved_file = await replicate.files.async_get(file_id)
+    else:
+        retrieved_file = replicate.files.get(file_id)
+
+    assert retrieved_file.id == file_id
+
+    # Test list
+    if async_flag:
+        file_list = await replicate.files.async_list()
+    else:
+        file_list = replicate.files.list()
+
+    assert file_list is not None
+    assert len(file_list) > 0
+    assert any(f.id == file_id for f in file_list)
+
+    # Test delete
+    if async_flag:
+        await replicate.files.async_delete(file_id)
+    else:
+        replicate.files.delete(file_id)
+
+    # Verify file is deleted
+    if async_flag:
+        file_list = await replicate.files.async_list()
+    else:
+        file_list = replicate.files.list()
+
+    assert all(f.id != file_id for f in file_list)