Add stream method on Prediction (#215)

Follow-up to #204

This PR adds a `stream` method on `Prediction`. In most cases, you'll
want to call `replicate.stream` to get the stream directly. But if you
have an existing prediction object and want to stream its output, you
can use this method instead. [^1]

[^1]: This rhymes with `run` being a shorthand for creating a prediction
and calling `wait`.

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

Author: mattt

replicate/prediction.py

@@ -5,11 +5,12 @@
 
 from typing_extensions import NotRequired, TypedDict, Unpack
 
-from replicate.exceptions import ModelError
+from replicate.exceptions import ModelError, ReplicateError
 from replicate.files import upload_file
 from replicate.json import encode_json
 from replicate.pagination import Page
 from replicate.resource import Namespace, Resource
+from replicate.stream import EventSource
 from replicate.version import Version
 
 try:
@@ -19,6 +20,7 @@
 
 if TYPE_CHECKING:
     from replicate.client import Client
+    from replicate.stream import ServerSentEvent
 
 
 class Prediction(Resource):
@@ -125,6 +127,25 @@ def wait(self) -> None:
             time.sleep(self._client.poll_interval)
             self.reload()
 
+    def stream(self) -> Optional[Iterator["ServerSentEvent"]]:
+        """
+        Stream the prediction output.
+
+        Raises:
+            ReplicateError: If the model does not support streaming.
+        """
+
+        url = self.urls and self.urls.get("stream", None)
+        if not url or not isinstance(url, str):
+            raise ReplicateError("Model does not support streaming")
+
+        headers = {}
+        headers["Accept"] = "text/event-stream"
+        headers["Cache-Control"] = "no-store"
+
+        with self._client._client.stream("GET", url, headers=headers) as response:
+            yield from EventSource(response)
+
     def cancel(self) -> None:
         """
         Cancels a running prediction.

tests/test_stream.py

@@ -32,3 +32,24 @@ async def test_stream(async_flag, record_mode):
 
     assert len(events) > 0
     assert events[0].event == "output"
+
+
+@pytest.mark.asyncio
+async def test_stream_prediction(record_mode):
+    if record_mode == "none":
+        return
+
+    version = "02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3"
+
+    input = {
+        "prompt": "Please write a haiku about llamas.",
+    }
+
+    prediction = replicate.predictions.create(version=version, input=input)
+
+    events = []
+    for event in prediction.stream():
+        events.append(event)
+
+    assert len(events) > 0
+    assert events[0].event == "output"