Refactor Stream module

This is now built on web standards based `EventSource` and
`ReadableStream` primatives. The latter is available in Node >= 18 and
the former is imported as a dependency from the `eventsource` module.

If available the implementation will use the native implementation
falling back to the module implementation which means this will work
in Browser environments and Deno.

Author: aron

index.js

@@ -1,6 +1,6 @@
 const ApiError = require("./lib/error");
 const ModelVersionIdentifier = require("./lib/identifier");
-const { Stream } = require("./lib/stream");
+const { createReadableStream } = require("./lib/stream");
 const {
   withAutomaticRetries,
   validateWebhook,
@@ -45,13 +45,17 @@ class Replicate {
    * @param {string} options.userAgent - Identifier of your app
    * @param {string} [options.baseUrl] - Defaults to https://api.replicate.com/v1
    * @param {Function} [options.fetch] - Fetch function to use. Defaults to `globalThis.fetch`
+   * @param {Function} [options.EventSource] - Custom EventSource implementation function to use.
+   * @param {Function} [options.ReadableStream] - Custom ReadableStream implementation function to use.
    */
   constructor(options = {}) {
     this.auth = options.auth || process.env.REPLICATE_API_TOKEN;
     this.userAgent =
       options.userAgent || `replicate-javascript/${packageJSON.version}`;
     this.baseUrl = options.baseUrl || "https://api.replicate.com/v1";
     this.fetch = options.fetch || globalThis.fetch;
+    this.EventSource = options.EventSource;
+    this.ReadableStream = options.ReadableStream;
 
     this.accounts = {
       current: accounts.current.bind(this),
@@ -289,9 +293,10 @@ class Replicate {
 
     if (prediction.urls && prediction.urls.stream) {
       const { signal } = options;
-      const stream = new Stream({
+      const stream = await createReadableStream({
         url: prediction.urls.stream,
-        fetch: this.fetch,
+        EventSource: this.EventSource,
+        ReadableStream: this.ReadableStream,
         options: { signal },
       });
       yield* stream;

index.test.ts

@@ -8,7 +8,7 @@ import Replicate, {
 } from "replicate";
 import nock from "nock";
 import fetch from "cross-fetch";
-import { Stream } from "./lib/stream";
+import { createReadableStream } from "./lib/stream";
 import { PassThrough } from "node:stream";
 
 let client: Replicate;
@@ -1188,19 +1188,24 @@ describe("Replicate client", () => {
 
   // Continue with tests for other methods
 
-  describe("Stream", () => {
-    function createStream(body: string | NodeJS.ReadableStream, status = 200) {
+  describe("createReadableStream", () => {
+    async function createStream(
+      body: string | NodeJS.ReadableStream,
+      status = 200
+    ) {
       const streamEndpoint = "https://stream.replicate.com";
       nock(streamEndpoint)
         .get("/fake_stream")
         .matchHeader("Accept", "text/event-stream")
         .reply(status, body);
 
-      return new Stream({ url: `${streamEndpoint}/fake_stream`, fetch });
+      return await createReadableStream({
+        url: `${streamEndpoint}/fake_stream`,
+      });
     }
 
     test("consumes a server sent event stream", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         event: output
         id: EVENT_1
@@ -1209,13 +1214,11 @@ describe("Replicate client", () => {
         event: done
         id: EVENT_2
         data: {}
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
-
       expect(await iterator.next()).toEqual({
         done: false,
         value: { event: "output", id: "EVENT_1", data: "hello world" },
@@ -1229,7 +1232,7 @@ describe("Replicate client", () => {
     });
 
     test("consumes multiple events", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         event: output
         id: EVENT_1
@@ -1242,9 +1245,8 @@ describe("Replicate client", () => {
         event: done
         id: EVENT_3
         data: {}
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
@@ -1266,7 +1268,7 @@ describe("Replicate client", () => {
     });
 
     test("ignores unexpected characters", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         : hi
 
@@ -1277,9 +1279,8 @@ describe("Replicate client", () => {
         event: done
         id: EVENT_2
         data: {}
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
@@ -1297,7 +1298,7 @@ describe("Replicate client", () => {
     });
 
     test("supports multiple lines of output in a single event", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         : hi
 
@@ -1310,9 +1311,8 @@ describe("Replicate client", () => {
         event: done
         id: EVENT_2
         data: {}
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
@@ -1335,7 +1335,7 @@ describe("Replicate client", () => {
 
     test("supports the server writing data lines in multiple chunks", async () => {
       const body = new PassThrough();
-      const stream = createStream(body);
+      const stream = await createStream(body);
 
       // Create a stream of data chunks split on the pipe character for readability.
       const data = `
@@ -1348,9 +1348,8 @@ describe("Replicate client", () => {
         event: done
         id: EVENT_2
         data: {}
-      `
-        .trim()
-        .replace(/^[ ]+/gm, "");
+
+      `.replace(/^[ ]+/gm, "");
 
       const chunks = data.split("|");
 
@@ -1392,7 +1391,7 @@ describe("Replicate client", () => {
 
     test("supports the server writing data in a complete mess", async () => {
       const body = new PassThrough();
-      const stream = createStream(body);
+      const stream = await createStream(body);
 
       // Create a stream of data chunks split on the pipe character for readability.
       const data = `
@@ -1407,9 +1406,8 @@ describe("Replicate client", () => {
         event: d|one
         id: EVENT|_2
         data: {}
-      `
-        .trim()
-        .replace(/^[ ]+/gm, "");
+
+      `.replace(/^[ ]+/gm, "");
 
       const chunks = data.split("|");
 
@@ -1450,15 +1448,13 @@ describe("Replicate client", () => {
     });
 
     test("supports ending without a done", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         event: output
         id: EVENT_1
         data: hello world
 
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
@@ -1470,7 +1466,7 @@ describe("Replicate client", () => {
     });
 
     test("an error event in the stream raises an exception", async () => {
-      const stream = createStream(
+      const stream = await createStream(
         `
         event: output
         id: EVENT_1
@@ -1480,27 +1476,23 @@ describe("Replicate client", () => {
         id: EVENT_2
         data: An unexpected error occurred
 
-      `
-          .trim()
-          .replace(/^[ ]+/gm, "")
+      `.replace(/^[ ]+/gm, "")
       );
 
       const iterator = stream[Symbol.asyncIterator]();
       expect(await iterator.next()).toEqual({
         done: false,
         value: { event: "output", id: "EVENT_1", data: "hello world" },
       });
-      await expect(iterator.next()).rejects.toThrowError(
-        "An unexpected error occurred"
-      );
+      await expect(iterator.next()).rejects.toThrowError("Unexpected Error");
       expect(await iterator.next()).toEqual({ done: true });
     });
 
     test("an error when fetching the stream raises an exception", async () => {
-      const stream = createStream("{}", 500);
+      const stream = await createStream("{}", 500);
       const iterator = stream[Symbol.asyncIterator]();
       await expect(iterator.next()).rejects.toThrowError(
-        "Request to https://stream.replicate.com/fake_stream failed with status 500 Internal Server Error: {}."
+        "Request to https://stream.replicate.com/fake_stream failed with status 500"
       );
       expect(await iterator.next()).toEqual({ done: true });
     });