Add replicate.stream method

Author: mattt

index.js

@@ -1,5 +1,6 @@
 const ApiError = require("./lib/error");
 const ModelVersionIdentifier = require("./lib/identifier");
+const { Stream } = require("./lib/stream");
 const { withAutomaticRetries } = require("./lib/util");
 
 const collections = require("./lib/collections");
@@ -235,6 +236,50 @@ class Replicate {
     return response;
   }
 
+  /**
+   * Stream a model and wait for its output.
+   *
+   * @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
+   * @param {object} [options.wait] - Options for waiting for the prediction to finish
+   * @param {number} [options.wait.interval] - Polling interval in milliseconds. Defaults to 250
+   * @param {number} [options.wait.max_attempts] - 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 {AbortSignal} [options.signal] - AbortSignal to cancel the prediction
+   * @throws {Error} If the prediction failed
+   * @yields {ServerSentEvent} Each streamed event from the prediction
+   */
+  async *stream(ref, options, progress) {
+    const { wait, ...data } = options;
+
+    const identifier = ModelVersionIdentifier.parse(ref);
+
+    let prediction;
+    if (identifier.version) {
+      prediction = await this.predictions.create({
+        ...data,
+        version: identifier.version,
+        stream: true,
+      });
+    } else {
+      prediction = await this.models.predictions.create(
+        identifier.owner,
+        identifier.name,
+        { ...data, stream: true }
+      );
+    }
+
+    if (prediction.urls && prediction.urls.stream) {
+      const { signal } = options;
+      const stream = new Stream(prediction.urls.stream, { signal });
+      yield* stream;
+    } else {
+      throw new Error("Prediction does not support streaming");
+    }
+  }
+
   /**
    * Paginate through a list of results.
    *

lib/stream.js

@@ -0,0 +1,100 @@
+const { Readable } = require("stream");
+
+class ServerSentEvent {
+  constructor(event, data, id, retry) {
+    this.event = event;
+    this.data = data;
+    this.id = id;
+    this.retry = retry;
+  }
+
+  toString() {
+    if (this.event === "output") {
+      return this.data;
+    }
+
+    return "";
+  }
+}
+
+class Stream extends Readable {
+  constructor(url, options) {
+    super();
+    this.url = url;
+    this.options = options;
+
+    this.event = null;
+    this.data = [];
+    this.lastEventId = null;
+    this.retry = null;
+  }
+
+  decode(line) {
+    if (!line) {
+      if (!this.event && !this.data.length && !this.lastEventId) {
+        return null;
+      }
+
+      const sse = new ServerSentEvent(
+        this.event,
+        this.data.join("\n"),
+        this.lastEventId
+      );
+
+      this.event = null;
+      this.data = [];
+      this.retry = null;
+
+      return sse;
+    }
+
+    if (line.startsWith(":")) {
+      return null;
+    }
+
+    const [field, value] = line.split(": ");
+    if (field === "event") {
+      this.event = value;
+    } else if (field === "data") {
+      this.data.push(value);
+    } else if (field === "id") {
+      this.lastEventId = value;
+    }
+
+    return null;
+  }
+
+  async *[Symbol.asyncIterator]() {
+    const response = await fetch(this.url, {
+      ...this.options,
+      headers: {
+        Accept: "text/event-stream",
+      },
+    });
+
+    for await (const chunk of response.body) {
+      const decoder = new TextDecoder("utf-8");
+      const text = decoder.decode(chunk);
+      const lines = text.split("\n");
+      for (const line of lines) {
+        const sse = this.decode(line);
+        if (sse) {
+          if (sse.event === "error") {
+            throw new Error(sse.data);
+          }
+
+          yield sse;
+
+          if (sse.event === "done") {
+            return;
+          }
+        }
+      }
+    }
+  }
+}
+
+module.exports = {
+  Stream,
+  ServerSentEvent,
+};