Overhaul encode_json logic to accomodate file encoding strategies

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

Author: mattt

replicate/deployment.py

@@ -1,11 +1,9 @@
-import asyncio
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, TypedDict, Union
 
 from typing_extensions import Unpack, deprecated
 
 from replicate.account import Account
-from replicate.file import base64_encode_file
-from replicate.json import encode_json
+from replicate.json import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.prediction import (
     Prediction,
@@ -424,9 +422,8 @@ def create(
         if input is not None:
             input = encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: self._client.files.create(file).urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 
@@ -449,13 +446,10 @@ async def async_create(
 
         file_encoding_strategy = params.pop("file_encoding_strategy", None)
         if input is not None:
-            input = encode_json(
+            input = await async_encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: asyncio.get_event_loop()
-                .run_until_complete(self._client.files.async_create(file))
-                .urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 
@@ -489,9 +483,8 @@ def create(
         if input is not None:
             input = encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: self._client.files.create(file).urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 
@@ -517,13 +510,10 @@ async def async_create(
 
         file_encoding_strategy = params.pop("file_encoding_strategy", None)
         if input is not None:
-            input = encode_json(
+            input = await async_encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: asyncio.get_event_loop()
-                .run_until_complete(self._client.files.async_create(file))
-                .urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 

replicate/file.py

@@ -1,15 +1,16 @@
-import base64
 import io
 import json
 import mimetypes
 import os
 import pathlib
 from typing import Any, BinaryIO, Dict, List, Optional, TypedDict, Union
 
-from typing_extensions import NotRequired, Unpack
+from typing_extensions import Literal, NotRequired, Unpack
 
 from replicate.resource import Namespace, Resource
 
+FileEncodingStrategy = Literal["base64", "url"]
+
 
 class File(Resource):
     """
@@ -168,26 +169,3 @@ def _create_file_params(
 
 def _json_to_file(json: Dict[str, Any]) -> File:  # pylint: disable=redefined-outer-name
     return File(**json)
-
-
-def base64_encode_file(file: io.IOBase) -> str:
-    """
-    Base64 encode a file.
-
-    Args:
-        file: A file handle to upload.
-    Returns:
-        str: A base64-encoded data URI.
-    """
-
-    file.seek(0)
-    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")
-
-    mime_type = (
-        mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
-    )
-    return f"data:{mime_type};base64,{encoded_body}"

replicate/json.py

@@ -1,7 +1,14 @@
+import base64
 import io
+import mimetypes
 from pathlib import Path
 from types import GeneratorType
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from replicate.client import Client
+    from replicate.file import FileEncodingStrategy
+
 
 try:
     import numpy as np  # type: ignore
@@ -14,22 +21,62 @@
 # pylint: disable=too-many-return-statements
 def encode_json(
     obj: Any,  # noqa: ANN401
-    upload_file: Callable[[io.IOBase], str],
+    client: "Client",
+    file_encoding_strategy: Optional["FileEncodingStrategy"] = None,
 ) -> Any:  # noqa: ANN401
     """
     Return a JSON-compatible version of the object.
     """
-    # Effectively the same thing as cog.json.encode_json.
 
     if isinstance(obj, dict):
-        return {key: encode_json(value, upload_file) for key, value in obj.items()}
+        return {
+            key: encode_json(value, client, file_encoding_strategy)
+            for key, value in obj.items()
+        }
+    if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
+        return [encode_json(value, client, file_encoding_strategy) for value in obj]
+    if isinstance(obj, Path):
+        with obj.open("rb") as file:
+            return encode_json(file, client, file_encoding_strategy)
+    if isinstance(obj, io.IOBase):
+        if file_encoding_strategy == "base64":
+            return base64.b64encode(obj.read()).decode("utf-8")
+        else:
+            return client.files.create(obj).urls["get"]
+    if HAS_NUMPY:
+        if isinstance(obj, np.integer):  # type: ignore
+            return int(obj)
+        if isinstance(obj, np.floating):  # type: ignore
+            return float(obj)
+        if isinstance(obj, np.ndarray):  # type: ignore
+            return obj.tolist()
+    return obj
+
+
+async def async_encode_json(
+    obj: Any,  # noqa: ANN401
+    client: "Client",
+    file_encoding_strategy: Optional["FileEncodingStrategy"] = None,
+) -> Any:  # noqa: ANN401
+    """
+    Asynchronously return a JSON-compatible version of the object.
+    """
+
+    if isinstance(obj, dict):
+        return {
+            key: (await async_encode_json(value, client, file_encoding_strategy))
+            for key, value in obj.items()
+        }
     if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
-        return [encode_json(value, upload_file) for value in obj]
+        return [
+            (await async_encode_json(value, client, file_encoding_strategy))
+            for value in obj
+        ]
     if isinstance(obj, Path):
         with obj.open("rb") as file:
-            return upload_file(file)
+            return encode_json(file, client, file_encoding_strategy)
     if isinstance(obj, io.IOBase):
-        return upload_file(obj)
+        return (await client.files.async_create(obj)).urls["get"]
     if HAS_NUMPY:
         if isinstance(obj, np.integer):  # type: ignore
             return int(obj)
@@ -38,3 +85,26 @@ def encode_json(
         if isinstance(obj, np.ndarray):  # type: ignore
             return obj.tolist()
     return obj
+
+
+def base64_encode_file(file: io.IOBase) -> str:
+    """
+    Base64 encode a file.
+
+    Args:
+        file: A file handle to upload.
+    Returns:
+        str: A base64-encoded data URI.
+    """
+
+    file.seek(0)
+    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")
+
+    mime_type = (
+        mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
+    )
+    return f"data:{mime_type};base64,{encoded_body}"

replicate/model.py

@@ -1,12 +1,10 @@
-import asyncio
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union, overload
 
 from typing_extensions import NotRequired, TypedDict, Unpack, deprecated
 
 from replicate.exceptions import ReplicateException
-from replicate.file import base64_encode_file
 from replicate.identifier import ModelVersionIdentifier
-from replicate.json import encode_json
+from replicate.json import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.prediction import (
     Prediction,
@@ -399,9 +397,8 @@ def create(
         if input is not None:
             input = encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: self._client.files.create(file).urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 
@@ -427,13 +424,10 @@ async def async_create(
 
         file_encoding_strategy = params.pop("file_encoding_strategy", None)
         if input is not None:
-            input = encode_json(
+            input = await async_encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: asyncio.get_event_loop()
-                .run_until_complete(self._client.files.async_create(file))
-                .urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(version=None, input=input, **params)
 

replicate/prediction.py

@@ -19,8 +19,8 @@
 from typing_extensions import NotRequired, TypedDict, Unpack
 
 from replicate.exceptions import ModelError, ReplicateError
-from replicate.file import base64_encode_file
-from replicate.json import encode_json
+from replicate.file import FileEncodingStrategy
+from replicate.json import async_encode_json, encode_json
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource
 from replicate.stream import EventSource
@@ -383,7 +383,7 @@ class CreatePredictionParams(TypedDict):
         stream: NotRequired[bool]
         """Enable streaming of prediction output."""
 
-        file_encoding_strategy: NotRequired[Literal["upload", "base64"]]
+        file_encoding_strategy: NotRequired[FileEncodingStrategy]
         """The strategy to use for encoding files in the prediction input."""
 
     @overload
@@ -460,9 +460,8 @@ def create(  # type: ignore
         if input is not None:
             input = encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: self._client.files.create(file).urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(
             version,
@@ -550,13 +549,10 @@ async def async_create(  # type: ignore
 
         file_encoding_strategy = params.pop("file_encoding_strategy", None)
         if input is not None:
-            input = encode_json(
+            input = await async_encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: asyncio.get_event_loop()
-                .run_until_complete(self._client.files.async_create(file))
-                .urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_prediction_body(
             version,

replicate/training.py

@@ -1,4 +1,3 @@
-import asyncio
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -14,9 +13,8 @@
 
 from typing_extensions import NotRequired, Unpack
 
-from replicate.file import base64_encode_file
 from replicate.identifier import ModelVersionIdentifier
-from replicate.json import encode_json
+from replicate.json import async_encode_json, encode_json
 from replicate.model import Model
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource
@@ -29,6 +27,7 @@
 
 if TYPE_CHECKING:
     from replicate.client import Client
+    from replicate.file import FileEncodingStrategy
 
 
 class Training(Resource):
@@ -221,7 +220,7 @@ class CreateTrainingParams(TypedDict):
         webhook: NotRequired[str]
         webhook_completed: NotRequired[str]
         webhook_events_filter: NotRequired[List[str]]
-        file_encoding_strategy: NotRequired[Literal["upload", "base64"]]
+        file_encoding_strategy: NotRequired["FileEncodingStrategy"]
 
     @overload
     def create(  # pylint: disable=too-many-arguments
@@ -283,10 +282,10 @@ def create(  # type: ignore
         if input is not None:
             input = encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: self._client.files.create(file).urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
+
         body = _create_training_body(input, **params)
 
         resp = self._client._request(
@@ -322,13 +321,10 @@ async def async_create(
 
         file_encoding_strategy = params.pop("file_encoding_strategy", None)
         if input is not None:
-            input = encode_json(
+            input = await async_encode_json(
                 input,
-                upload_file=base64_encode_file
-                if file_encoding_strategy == "base64"
-                else lambda file: asyncio.get_event_loop()
-                .run_until_complete(self._client.files.async_create(file))
-                .urls["get"],
+                client=self._client,
+                file_encoding_strategy=file_encoding_strategy,
             )
         body = _create_training_body(input, **params)