Chain predictions off versions

This connects predictions to versions conceptually, to match their
concept elsewhere, and gives us a clear API for differentiating from
other kinds of predictions in the future.

We create `Version` and `Prediction` classes we don't expect to expose
externally. The entrypoint for users continues to be methods called from
`ReplicateClient` directly.

Author: erbridge

README.md

@@ -15,12 +15,9 @@ To run a prediction and return its output:
 ```js
 import replicate from "replicate";
 
-// Set your model's input parameters here
-const input = {
+const prediction = await replicate.version("<MODEL VERSION>").predict({
   prompt: "painting of a cat by andy warhol",
-};
-
-const prediction = await replicate.predict("<MODEL VERSION>", input);
+});
 
 console.log(prediction.output);
 // "https://replicate.delivery/pbxt/lGWovsQZ7jZuNtPvofMth1rSeCcVn5xes8dWWdWZ64MlTi7gA/out-0.png"
@@ -32,18 +29,34 @@ running, you can pass in an `onUpdate` callback function:
 ```js
 import replicate from "replicate";
 
-// Set your model's input parameters here
-const input = {
-  prompt: "painting of a cat by andy warhol",
-};
-
-await replicate.predict("<MODEL VERSION>", input, {
-  onUpdate: (prediction) => {
-    console.log(prediction.output);
+await replicate.version("<MODEL VERSION>").predict(
+  {
+    prompt: "painting of a cat by andy warhol",
   },
+  {
+    onUpdate: (prediction) => {
+      console.log(prediction.output);
+    },
+  }
+);
+```
+
+If you'd prefer to control your own polling you can use the low-level
+`createPrediction()` method:
+
+```js
+import replicate from "replicate";
+
+const prediction = await replicate.version("<MODEL VERSION>").createPrediction({
+  prompt: "painting of a cat by andy warhol",
 });
+
+console.log(prediction.status); // "starting"
 ```
 
+From there, you can fetch the current status of the prediction using
+`await prediction.load()` or `await replicate.prediction(prediction.id).load()`.
+
 ## License
 
 [Apache 2.0](LICENSE)

lib/Prediction.js

@@ -1,3 +1,6 @@
+import { ReplicateError } from "./errors.js";
+import ReplicateObject from "./ReplicateObject.js";
+
 export const PredictionStatus = {
   STARTING: "starting",
   PROCESSING: "processing",
@@ -6,28 +9,19 @@ export const PredictionStatus = {
   FAILED: "failed",
 };
 
-export default class Prediction {
-  static fromJSON(json) {
-    let props;
-    try {
-      props = JSON.parse(json);
-    } catch {
-      throw new Error(`Unable to parse JSON: ${json}`);
-    }
-
-    return new this(props);
-  }
+export default class Prediction extends ReplicateObject {
+  constructor({ id, ...rest }, client) {
+    super(rest, client);
 
-  id;
-  status;
+    if (!id) {
+      throw new ReplicateError("id is required");
+    }
 
-  constructor({ id, status, ...rest }) {
     this.id = id;
-    this.status = status;
+  }
 
-    for (const key in rest) {
-      this[key] = rest[key];
-    }
+  actionForGet() {
+    return `GET /v1/predictions/${this.id}`;
   }
 
   hasTerminalStatus() {

lib/Prediction.test.js

@@ -0,0 +1,57 @@
+import { jest } from "@jest/globals";
+import { FetchError } from "node-fetch";
+import Prediction, { PredictionStatus } from "./Prediction.js";
+
+jest.unstable_mockModule("node-fetch", () => ({
+  default: jest.fn(),
+  FetchError,
+}));
+
+const { default: ReplicateClient } = await import("./ReplicateClient.js");
+
+let client;
+let prediction;
+
+beforeEach(() => {
+  process.env.REPLICATE_API_TOKEN = "test-token-from-env";
+
+  client = new ReplicateClient({});
+  prediction = client.prediction("test-prediction");
+});
+
+describe("load()", () => {
+  it("makes request to get prediction", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "test-prediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    await prediction.load();
+
+    expect(client.request).toHaveBeenCalledWith(
+      "GET /v1/predictions/test-prediction"
+    );
+  });
+
+  it("returns a Prediction", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "test-prediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    const returnedPrediction = await prediction.load();
+
+    expect(returnedPrediction).toBeInstanceOf(Prediction);
+  });
+
+  it("updates the prediction in place", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "test-prediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    const returnedPrediction = await prediction.load();
+
+    expect(returnedPrediction).toBe(prediction);
+  });
+});

lib/ReplicateClient.js

@@ -6,7 +6,7 @@ import {
   ReplicateResponseError,
 } from "./errors.js";
 import Prediction from "./Prediction.js";
-import { sleep } from "./utils.js";
+import Version from "./Version.js";
 
 export default class ReplicateClient {
   baseURL;
@@ -21,79 +21,12 @@ export default class ReplicateClient {
     }
   }
 
-  // TODO: Optionally autocancel prediction on exception.
-  async predict(
-    version,
-    input,
-    { onUpdate, onTemporaryError } = {},
-    {
-      defaultPollingInterval = 500,
-      backoffFn = (errorCount) => Math.pow(2, errorCount) * 100,
-    } = {}
-  ) {
-    if (!version) {
-      throw new ReplicateError("version is required");
-    }
-
-    if (!input) {
-      throw new ReplicateError("input is required");
-    }
-
-    let prediction = await this.createPrediction(version, input);
-
-    onUpdate && onUpdate(prediction);
-
-    let pollingInterval = defaultPollingInterval;
-    let errorCount = 0;
-
-    while (!prediction.hasTerminalStatus()) {
-      await sleep(pollingInterval);
-      pollingInterval = defaultPollingInterval; // Reset to default each time.
-
-      try {
-        prediction = await this.getPrediction(prediction.id);
-
-        onUpdate && onUpdate(prediction);
-
-        errorCount = 0; // Reset because we've had a non-error response.
-      } catch (err) {
-        if (!err instanceof ReplicateResponseError) {
-          throw err;
-        }
-
-        if (
-          !err.status ||
-          (Math.floor(err.status / 100) !== 5 && err.status !== 429)
-        ) {
-          throw err;
-        }
-
-        errorCount += 1;
-
-        onTemporaryError && onTemporaryError(err);
-
-        pollingInterval = backoffFn(errorCount);
-      }
-    }
-
-    return prediction;
-  }
-
-  async createPrediction(version, input) {
-    const predictionData = await this.request("POST /v1/predictions", {
-      version,
-      input,
-    });
-
-    return new Prediction(predictionData);
+  version(id) {
+    return new Version({ id }, this);
   }
 
-  async getPrediction(predictionID) {
-    const predictionData = await this.request(
-      `GET /v1/predictions/${predictionID}`
-    );
-
-    return new Prediction(predictionData);
+  prediction(id) {
+    return new Prediction({ id }, this);
   }
 
   async request(action, body) {