Implement Stream interface with prediction property

Author: mattt

index.d.ts

@@ -8,6 +8,16 @@ declare module 'replicate' {
     results: T[];
   }
 
+  interface Stream {
+    response: Response;
+    reader: ReadableStream;
+    decoder: TextDecoder;
+    output: string;
+    prediction?: Promise<Prediction>;
+    next(): Promise<{ value: string, done: boolean } | { done: true }>;
+    [ Symbol.asyncIterator ](): this;
+  }
+
   export interface Collection {
     name: string;
     slug: string;
@@ -89,7 +99,7 @@ declare module 'replicate' {
     fetch: Function;
 
 
-    run(identifier: RunIdentifier, options: RunOptions & { stream: true }): Promise<AsyncIterable<any>>;
+    run(identifier: RunIdentifier, options: RunOptions & { stream: true }): Promise<Stream>;
     run(identifier: RunIdentifier, options: RunOptions): Promise<object>;
 
     request(route: string, parameters: any): Promise<any>;
@@ -120,7 +130,7 @@ declare module 'replicate' {
     };
 
     predictions: {
-      create(options: CreatePredictionOptions & { stream: true }): Promise<AsyncIterable<any>>;
+      create(options: CreatePredictionOptions & { stream: true }): Promise<Stream>;
       create(options: CreatePredictionOptions): Promise<Prediction>;
       get(prediction_id: string): Promise<Prediction>;
       cancel(prediction_id: string): Promise<Prediction>;

index.js

@@ -3,7 +3,6 @@ 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
@@ -138,16 +137,27 @@ class Replicate {
    * @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
+   * @returns {Promise<Response>} - Resolves with the response object
    * @throws {Error} If the request failed
    */
   async request(route, parameters) {
     const { auth, baseUrl, userAgent } = this;
 
-    const url = new URL(
-      route.startsWith('/') ? route.slice(1) : route,
-      baseUrl.endsWith('/') ? baseUrl : `${baseUrl}/`
-    );
+    let url;
+    if (route instanceof URL) {
+      if (route.origin !== baseUrl) {
+        throw new Error(
+          `Invalid URL: ${route.origin} does not match ${baseUrl}`
+        );
+      }
+
+      url = route;
+    } else {
+      url = new URL(
+        route.startsWith('/') ? route.slice(1) : route,
+        baseUrl.endsWith('/') ? baseUrl : `${baseUrl}/`
+      );
+    }
 
     const { method = 'GET', params = {}, data } = parameters;
 
@@ -175,12 +185,7 @@ 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();
+    return response;
   }
 
   /**
@@ -198,7 +203,7 @@ class Replicate {
     const response = await endpoint();
     yield response.results;
     if (response.next) {
-      const nextPage = () => this.request(response.next, { method: 'GET' });
+      const nextPage = () => this.request(response.next, { method: 'GET' }).then((r) => r.json());
       yield* this.paginate(nextPage);
     }
   }

index.test.ts

@@ -518,7 +518,38 @@ describe('Replicate client', () => {
         [
           rest.post('https://api.replicate.com/v1/predictions/', (req, res, ctx) => {
             expect(req.headers.get('Accept')).toBe('text/event-stream');
-            return res(ctx.status(201), ctx.text('data: Once upon a time\n\n'));
+            return res(
+              ctx.status(201),
+              ctx.set('Content-Type', 'text/event-stream'),
+              ctx.set('Link', '<https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq>; rel="self"'),
+              ctx.body('data: Once upon a time\n\n'),
+            );
+          }),
+          rest.get('https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq', (req, res, ctx) => {
+            return res(
+              ctx.status(200),
+              ctx.json({
+                id: 'ufawqhfynnddngldkgtslldrkq',
+                version:
+                  '5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa',
+                urls: {
+                  get: 'https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq',
+                  cancel:
+                    'https://api.replicate.com/v1/predictions/ufawqhfynnddngldkgtslldrkq/cancel',
+                },
+                created_at: '2022-04-26T22:13:06.224088Z',
+                started_at: null,
+                completed_at: null,
+                status: 'starting',
+                input: {
+                  text: 'Alice',
+                },
+                output: null,
+                error: null,
+                logs: null,
+                metrics: {},
+              })
+            );
           }),
         ],
         async () => {
@@ -530,6 +561,10 @@ describe('Replicate client', () => {
             }
           );
 
+          const prediction = await stream.prediction;
+          expect(prediction.id).toBe('ufawqhfynnddngldkgtslldrkq');
+          expect(prediction.status).toBe('starting');
+
           for await (const event of await stream) {
             expect(event).toBe('Once upon a time');
           }

lib/collections.js

@@ -5,9 +5,11 @@
  * @returns {Promise<object>} - Resolves with the collection data
  */
 async function getCollection(collection_slug) {
-  return this.request(`/collections/${collection_slug}`, {
+  const response = await this.request(`/collections/${collection_slug}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -16,9 +18,11 @@ async function getCollection(collection_slug) {
  * @returns {Promise<object>} - Resolves with the collections data
  */
 async function listCollections() {
-  return this.request('/collections', {
+  const response = await this.request('/collections', {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 module.exports = { get: getCollection, list: listCollections };

lib/models.js

@@ -6,9 +6,11 @@
  * @returns {Promise<object>} Resolves with the model data
  */
 async function getModel(model_owner, model_name) {
-  return this.request(`/models/${model_owner}/${model_name}`, {
+  const response = await this.request(`/models/${model_owner}/${model_name}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -19,9 +21,11 @@ async function getModel(model_owner, model_name) {
  * @returns {Promise<object>} Resolves with the list of model versions
  */
 async function listModelVersions(model_owner, model_name) {
-  return this.request(`/models/${model_owner}/${model_name}/versions`, {
+  const response = await this.request(`/models/${model_owner}/${model_name}/versions`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -33,12 +37,14 @@ async function listModelVersions(model_owner, model_name) {
  * @returns {Promise<object>} Resolves with the model version data
  */
 async function getModelVersion(model_owner, model_name, version_id) {
-  return this.request(
+  const response = await this.request(
     `/models/${model_owner}/${model_name}/versions/${version_id}`,
     {
       method: 'GET',
     }
   );
+
+  return response.json();
 }
 
 module.exports = {

lib/predictions.js

@@ -1,7 +1,8 @@
+const { Stream } = require('./stream');
+
 /**
  * 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
@@ -11,7 +12,7 @@
  * @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
- * @returns {Promise<object>|AsyncIterable<any>} Resolves with the created prediction, or a stream of prediction output if `options.stream` is `true`
+ * @returns {Promise<object>|Stream} Resolves with the created prediction, or a stream of prediction output if `options.stream` is `true`
  */
 async function createPrediction(options) {
   const { wait, stream, ...data } = options;
@@ -34,23 +35,49 @@ async function createPrediction(options) {
     headers.append('Accept', 'text/event-stream');
   }
 
-  const prediction = this.request('/predictions', {
+  const response = await this.request('/predictions', {
     method: 'POST',
     data,
     headers,
   });
 
+  if (stream) {
+    const contentType = response.headers.get('content-type');
+    if (contentType && contentType.startsWith('text/event-stream')) {
+      const sse = Stream(response);
+
+      // Get the prediction ID from the Link header
+      const link = response.headers.get('Link');
+      const match = link.match(/<(.*)>; rel="self"/);
+      if (match && match[1]) {
+        const ref = response.headers.get('Link').match(/<(.*)>; rel="self"/)[1];
+        try {
+          const url = new URL(ref);
+          const prediction_id = url.pathname.split('/').pop();
+
+          if (prediction_id) {
+            Object.defineProperty(sse, 'prediction', async () => (this.getPrediction(prediction_id)));
+          }
+        } catch (err) {
+          throw new Error('Unable to get prediction ID from response');
+        }
+      }
+
+      return sse;
+    }
+  }
+
   if (wait) {
     const { maxAttempts, interval } = wait;
 
     // eslint-disable-next-line no-promise-executor-return
     const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
     await sleep(interval || 250);
 
-    return this.wait(await prediction, { maxAttempts, interval });
+    return this.wait(await response.json(), { maxAttempts, interval });
   }
 
-  return prediction;
+  return response.json();
 }
 
 /**
@@ -60,9 +87,11 @@ async function createPrediction(options) {
  * @returns {Promise<object>} Resolves with the prediction data
  */
 async function getPrediction(prediction_id) {
-  return this.request(`/predictions/${prediction_id}`, {
+  const response = await this.request(`/predictions/${prediction_id}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -72,9 +101,11 @@ async function getPrediction(prediction_id) {
  * @returns {Promise<object>} Resolves with the data for the training
  */
 async function cancelPrediction(prediction_id) {
-  return this.request(`/predictions/${prediction_id}/cancel`, {
+  const response = await this.request(`/predictions/${prediction_id}/cancel`, {
     method: 'POST',
   });
+
+  return response.json();
 }
 
 /**
@@ -83,9 +114,11 @@ async function cancelPrediction(prediction_id) {
  * @returns {Promise<object>} - Resolves with a page of predictions
  */
 async function listPredictions() {
-  return this.request('/predictions', {
+  const response = await this.request('/predictions', {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 module.exports = {

lib/stream.js

@@ -0,0 +1,41 @@
+/**
+ * An async iterable that yields Server-Sent Event messages from a response.
+ */
+class Stream {
+  /**
+   * Create a new Server-Sent Event stream.
+   *
+   * @param {Response} response - The server response
+   */
+  constructor(response) {
+    this.response = response;
+    this.reader = response.body.getReader();
+    this.decoder = new TextDecoder('utf-8');
+    this.output = '';
+  }
+
+  async next() {
+    let eventEndIndex;
+    while ((eventEndIndex = this.output.indexOf('\n\n')) < 0) { // eslint-disable-line no-cond-assign
+      const { done, chunk } = await this.reader.read(); // eslint-disable-line no-await-in-loop
+      if (done) {
+        return { done: true };
+      }
+      this.output += this.decoder.decode(chunk, { stream: true });
+    }
+
+    const event = this.output.slice(0, eventEndIndex);
+    this.output = this.output.slice(eventEndIndex + 2);
+
+    const message = this.parseEvent(event);
+    return { value: message, done: false };
+  }
+
+  [Symbol.asyncIterator]() {
+    return this;
+  }
+}
+
+module.exports = {
+  Stream
+};