Add support to FileOutput to read data-uris

Author: aron

replicate/helpers.py

@@ -131,21 +131,37 @@ def __init__(self, url: str, client: "Client") -> None:
         self._client = client
 
     def read(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(',', 1)
+            return base64.b64decode(encoded)
+
         with self._client._client.stream("GET", self.url) as response:
             response.raise_for_status()
             return response.read()
 
     def __iter__(self) -> Iterator[bytes]:
+        if self.url.startswith("data:"):
+            yield self.read()
+            return
+
         with self._client._client.stream("GET", self.url) as response:
             response.raise_for_status()
             yield from response.iter_bytes()
 
     async def aread(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(',', 1)
+            return base64.b64decode(encoded)
+
         async with self._client._async_client.stream("GET", self.url) as response:
             response.raise_for_status()
             return await response.aread()
 
     async def __aiter__(self) -> AsyncIterator[bytes]:
+        if self.url.startswith("data:"):
+            yield await self.aread()
+            return
+
         async with self._client._async_client.stream("GET", self.url) as response:
             response.raise_for_status()
             async for chunk in response.aiter_bytes():

tests/test_run.py

@@ -426,3 +426,84 @@ def prediction_with_status(
 
     assert output1.read() == b"Hello,"
     assert output2.read() == b" world!"
+
+
+@pytest.mark.asyncio
+async def test_run_with_file_output_data_uri(mock_replicate_api_token):
+    def prediction_with_status(
+        status: str, output: str | list[str] | None = None
+    ) -> dict:
+        return {
+            "id": "p1",
+            "model": "test/example",
+            "version": "v1",
+            "urls": {
+                "get": "https://api.replicate.com/v1/predictions/p1",
+                "cancel": "https://api.replicate.com/v1/predictions/p1/cancel",
+            },
+            "created_at": "2023-10-05T12:00:00.000000Z",
+            "source": "api",
+            "status": status,
+            "input": {"text": "world"},
+            "output": output,
+            "error": "OOM" if status == "failed" else None,
+            "logs": "",
+        }
+
+    router = respx.Router(base_url="https://api.replicate.com/v1")
+    router.route(method="POST", path="/predictions").mock(
+        return_value=httpx.Response(
+            201,
+            json=prediction_with_status("processing"),
+        )
+    )
+    router.route(method="GET", path="/predictions/p1").mock(
+        return_value=httpx.Response(
+            200,
+            json=prediction_with_status(
+                "succeeded",
+                "data:text/plain;base64,SGVsbG8sIHdvcmxkIQ==",
+            ),
+        )
+    )
+    router.route(
+        method="GET",
+        path="/models/test/example/versions/v1",
+    ).mock(
+        return_value=httpx.Response(
+            201,
+            json={
+                "id": "f2d6b24e6002f25f77ae89c2b0a5987daa6d0bf751b858b94b8416e8542434d1",
+                "created_at": "2024-07-18T00:35:56.210272Z",
+                "cog_version": "0.9.10",
+                "openapi_schema": {
+                    "openapi": "3.0.2",
+                },
+            },
+        )
+    )
+
+    client = Client(
+        api_token="test-token", transport=httpx.MockTransport(router.handler)
+    )
+    client.poll_interval = 0.001
+
+    output = cast(
+        FileOutput,
+        client.run(
+            "test/example:v1",
+            input={
+                "text": "Hello, world!",
+            },
+            use_file_output=True,
+        ),
+    )
+
+    assert output.url == "data:text/plain;base64,SGVsbG8sIHdvcmxkIQ=="
+    assert output.read() == b"Hello, world!"
+    for chunk in output:
+        assert chunk == b"Hello, world!"
+
+    assert await output.aread() == b"Hello, world!"
+    async for chunk in output:
+        assert chunk == b"Hello, world!"