Implement streaming responses for run and predictions.create endpoints

Author: mattt

index.d.ts

@@ -58,6 +58,21 @@ declare module 'replicate' {
 
   export type Training = Prediction;
 
+  interface RunOptions {
+    input: any;
+    wait?: {
+      interval?: number;
+      maxAttempts?: number;
+    };
+    webhook?: string;
+    webhook_events_filter?: string[];
+    stream?: boolean;
+  }
+
+  interface CreatePredictionOptions extends RunOptions {
+    version: string;
+  }
+
   export default class Replicate {
     constructor(options: {
       auth: string;
@@ -71,15 +86,10 @@ declare module 'replicate' {
     baseUrl?: string;
     fetch: Function;
 
-    run(
-      identifier: `${string}/${string}:${string}`,
-      options: {
-        input: object;
-        wait?: boolean | { interval?: number; maxAttempts?: number };
-        webhook?: string;
-        webhook_events_filter?: WebhookEventType[];
-      }
-    ): Promise<object>;
+
+    run(identifier: string, options: RunOptions & { stream: true }): Promise<AsyncIterable<any>>;
+    run(identifier: string, options: RunOptions): Promise<object>;
+
     request(route: string, parameters: any): Promise<any>;
     paginate<T>(endpoint: () => Promise<Page<T>>): AsyncGenerator<[ T ]>;
     wait(
@@ -108,12 +118,8 @@ declare module 'replicate' {
     };
 
     predictions: {
-      create(options: {
-        version: string;
-        input: object;
-        webhook?: string;
-        webhook_events_filter?: WebhookEventType[];
-      }): Promise<Prediction>;
+      create(options: CreatePredictionOptions & { stream: true }): Promise<AsyncIterable<any>>;
+      create(options: CreatePredictionOptions): Promise<Prediction>;
       get(prediction_id: string): Promise<Prediction>;
       cancel(prediction_id: string): Promise<Prediction>;
       list(): Promise<Page<Prediction>>;

index.js

@@ -3,6 +3,7 @@ const models = require('./lib/models');
 const predictions = require('./lib/predictions');
 const trainings = require('./lib/trainings');
 const packageJSON = require('./package.json');
+const { handleSSE } = require('./lib/util');
 
 /**
  * Replicate API client library
@@ -74,6 +75,7 @@ class Replicate {
   /**
    * Run a model and wait for its output.
    *
+   * @typedef AsyncIterable
    * @param {string} identifier - Required. The model version identifier in the format "{owner}/{name}:{version}"
    * @param {object} options
    * @param {object} options.input - Required. An object with the model inputs
@@ -82,8 +84,9 @@ class Replicate {
    * @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`)
+   * @param {boolean} [options.stream] - Whether to stream the prediction output. Defaults to false
    * @throws {Error} If the prediction failed
-   * @returns {Promise<object>} - Resolves with the output of running the model
+   * @returns {Promise<object>|Promise<AsyncIterable<any>>} - Resolves with the output of running the model, or a stream of output objects if `options.stream` is `true`
    */
   async run(identifier, options) {
     // Define a pattern for owner and model names that allows
@@ -105,6 +108,14 @@ class Replicate {
     }
 
     const { version } = match.groups;
+
+    if (options.stream) {
+      return this.predictions.create({
+        ...options,
+        version,
+      });
+    }
+
     const prediction = await this.predictions.create({
       wait: true,
       ...options,
@@ -125,6 +136,7 @@ class Replicate {
    * @param {object} parameters - Request parameters
    * @param {string} [parameters.method] - HTTP method. Defaults to GET
    * @param {object} [parameters.params] - Query parameters
+   * @param {object|Headers} [parameters.headers] - HTTP headers
    * @param {object} [parameters.data] - Body parameters
    * @returns {Promise<object>} - Resolves with the API response data
    * @throws {Error} If the request failed
@@ -143,11 +155,15 @@ class Replicate {
       url.searchParams.append(key, value);
     });
 
-    const headers = {
-      Authorization: `Token ${auth}`,
-      'Content-Type': 'application/json',
-      'User-Agent': userAgent,
-    };
+    const headers = new Headers();
+    headers.append('Authorization', `Token ${auth}`);
+    headers.append('Content-Type', 'application/json');
+    headers.append('User-Agent', userAgent);
+    if (parameters.headers) {
+      parameters.headers.forEach((value, key) => {
+        headers.append(key, value);
+      });
+    }
 
     const response = await this.fetch(url, {
       method,
@@ -159,6 +175,11 @@ class Replicate {
       throw new Error(`API request failed: ${response.statusText}`);
     }
 
+    const contentType = response.headers.get('content-type');
+    if (contentType && contentType.startsWith('text/event-stream')) {
+      return handleSSE(response);
+    }
+
     return response.json();
   }
 

lib/predictions.js

@@ -1,6 +1,7 @@
 /**
  * Create a new prediction
  *
+ * @typedef AsyncIterable
  * @param {object} options
  * @param {string} options.version - Required. The model version
  * @param {object} options.input - Required. An object with the model inputs
@@ -9,10 +10,15 @@
  * @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>|AsyncIterable<any>} Resolves with the created prediction, or a stream of prediction output if `options.stream` is `true`
  */
 async function createPrediction(options) {
-  const { wait, ...data } = options;
+  const { wait, stream, ...data } = options;
+
+  if (stream && wait) {
+    throw new Error('Incompatible options: stream and wait');
+  }
 
   if (data.webhook) {
     try {
@@ -23,9 +29,15 @@ async function createPrediction(options) {
     }
   }
 
+  const headers = new Headers();
+  if (stream) {
+    headers.append('Accept', 'text/event-stream');
+  }
+
   const prediction = this.request('/predictions', {
     method: 'POST',
     data,
+    headers,
   });
 
   if (wait) {

lib/util.js

@@ -0,0 +1,41 @@
+/**
+ * An async iterable that yields Server-Sent Event messages from a response.
+ * @param {Response} response - SSE response
+ * @yields {object} - SSE message
+ */
+async function* handleSSE(response) {
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder('utf-8');
+  let data = '';
+
+  for (; ;) {
+    const { done, chunk } = await reader.read(); // eslint-disable-line no-await-in-loop
+    if (done) {
+      break;
+    }
+
+    data += decoder.decode(chunk, { stream: true });
+
+    let eventEndIndex;
+    // eslint-disable-next-line no-cond-assign
+    while ((eventEndIndex = data.indexOf('\n\n')) >= 0) {
+      const event = data.slice(0, eventEndIndex);
+      data = data.slice(eventEndIndex + 2);
+
+      const message = {};
+      event.split('\n').forEach((line) => {
+        const [field, ...rest] = line.split(':');
+        const value = rest.join(':').trim();
+        if (field) {
+          message[field] = value;
+        }
+      });
+
+      yield message;
+    }
+  }
+}
+
+module.exports = {
+  handleSSE
+};