Add stream parameter to predictions.create

Author: mattt

README.md

@@ -251,17 +251,14 @@ const response = await replicate.predictions.create(options);
 | ------------------------------- | -------- | -------------------------------------------------------------------------------------------------------------------------------- |
 | `options.version`               | string   | **Required**. The model version                                                                                                  |
 | `options.input`                 | object   | **Required**. An object with the model's inputs                                                                                  |
+| `options.stream`                | boolean  | Requests a URL for streaming output output                                                                                       |
 | `options.webhook`               | string   | An HTTPS URL for receiving a webhook when the prediction has new output                                                          |
 | `options.webhook_events_filter` | string[] | You can change which events trigger webhook requests by specifying webhook events (`start` \| `output` \| `logs` \| `completed`) |
 
 ```jsonc
 {
   "id": "ufawqhfynnddngldkgtslldrkq",
   "version": "5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-  "urls": {
-    "get": "https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq",
-    "cancel": "https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq/cancel"
-  },
   "status": "succeeded",
   "input": {
     "text": "Alice"
@@ -272,10 +269,56 @@ const response = await replicate.predictions.create(options);
   "metrics": {},
   "created_at": "2022-04-26T22:13:06.224088Z",
   "started_at": null,
-  "completed_at": null
+  "completed_at": null,
+  "urls": {
+    "get": "https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq",
+    "cancel": "https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq/cancel",
+    "stream": "https://streaming.api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq" // Present only if `options.stream` is `true`
+  }
 }
 ```
 
+#### Streaming
+
+Specify the `stream` option when creating a prediction
+to request a URL to receive streaming output using
+[server-sent events (SSE)](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events).
+
+If the requested model version supports streaming,
+then the returned prediction will have a `stream` entry in its `urls` property
+with a URL that you can use to construct an
+[`EventSource`](https://developer.mozilla.org/en-US/docs/Web/API/EventSource).
+
+```js
+if (prediction && prediction.urls && prediction.urls.stream) {
+  const source = new EventSource(prediction.urls.stream);
+
+  source.addEventListener("output", (e) => {
+    console.log("output", e.data);
+  });
+
+  source.addEventListener("error"), (e) => {
+    console.error("error", JSON.parse(e.data));
+  });
+
+  source.addEventListener("done"), (e) => {
+    source.close();
+    console.log("done", JSON.parse(e.data));
+  });
+}
+```
+
+A prediction's event stream consists of the following event types:
+
+| event    | format     | description                                    |
+| -------- | ---------- | ---------------------------------------------- |
+| `output` | plain text | Emitted when the prediction returns new output |
+| `error`  | JSON       | Emitted when the prediction returns an error   |
+| `done`   | JSON       | Emitted when the prediction finishes           |
+
+A `done` event is emitted when a prediction finishes successfully,
+is cancelled, or produces an error.
+
 ### `replicate.predictions.get`
 
 ```js

index.d.ts

@@ -53,6 +53,11 @@ declare module 'replicate' {
     created_at: string;
     updated_at: string;
     completed_at?: string;
+    urls: {
+      get: string;
+      cancel: string;
+      stream?: string;
+    };
   }
 
   export type Training = Prediction;
@@ -123,6 +128,7 @@ declare module 'replicate' {
       create(options: {
         version: string;
         input: object;
+        stream?: boolean;
         webhook?: string;
         webhook_events_filter?: WebhookEventType[];
       }): Promise<Prediction>;

index.test.ts

@@ -152,6 +152,24 @@ describe('Replicate client', () => {
       expect(prediction.id).toBe('ufawqhfynnddngldkgtslldrkq');
     });
 
+    test('Passes stream parameter to API endpoint', async () => {
+      nock(BASE_URL)
+        .post('/predictions')
+        .reply(201, (_uri, body) => {
+          expect(body[ 'stream' ]).toBe(true);
+          return body
+        })
+
+      await client.predictions.create({
+        version:
+          '5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa',
+        input: {
+          prompt: 'Tell me a story',
+        },
+        stream: true
+      });
+    });
+
     test('Throws an error if webhook URL is invalid', async () => {
       await expect(async () => {
         await client.predictions.create({
@@ -506,7 +524,7 @@ describe('Replicate client', () => {
           status: 'processing',
         })
         .get('/predictions/ufawqhfynnddngldkgtslldrkq')
-        .reply(200, {
+        .reply(201, {
           id: 'ufawqhfynnddngldkgtslldrkq',
           status: 'succeeded',
           output: 'foobar',

lib/predictions.js

@@ -9,10 +9,11 @@
  * @param {number} [options.wait.maxAttempts] - Maximum number of polling attempts. Defaults to no limit
  * @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the prediction has new output
  * @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`)
- * @returns {Promise<object>} Resolves with the created prediction data
+ * @param {boolean} [options.stream] - Whether to stream the prediction output. Defaults to false
+ * @returns {Promise<object>} Resolves with the created prediction
  */
 async function createPrediction(options) {
-  const { wait, ...data } = options;
+  const { stream, ...data } = options;
 
   if (data.webhook) {
     try {