Support webhooks when creating predictions

Author: erbridge

README.md

@@ -75,6 +75,33 @@ 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()`.
 
+You can also provide webhook configuration to have Replicate send POST requests
+to your service when certain events occur:
+
+```js
+import replicate from "replicate";
+
+await replicate
+  .model(
+    "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf"
+  )
+  .createPrediction(
+    {
+      prompt: "painting of a cat by andy warhol",
+    },
+    {
+      webhook: {
+        // See https://replicate.com/docs/reference/http#create-prediction--webhook
+        url: "https://your.host/webhook",
+
+        // Optional
+        // See https://replicate.com/docs/reference/http#create-prediction--webhook_events_filter
+        events: ["output", "completed"],
+      },
+    }
+  );
+```
+
 ## Contributing
 
 While we'd love to accept contributions to this library, please open an issue

lib/Model.js

@@ -122,15 +122,27 @@ export default class Model extends ReplicateObject {
     return prediction;
   }
 
-  async createPrediction(input) {
+  async createPrediction(input, { webhook = {} } = {}) {
     // This is here and not on `Prediction` because conceptually, a prediction
     // from a model "belongs" to the model. It's an odd feature of the API that
     // the prediction creation isn't an action on the model (or that it doesn't
     // actually use the model information, only the version), but we don't need
     // to expose that to users of this library.
+
+    let webhookUrl;
+    let webhookEvents;
+    if (typeof webhook === "string") {
+      webhookUrl = webhook;
+    } else {
+      webhookUrl = webhook.url;
+      webhookEvents = webhook.events;
+    }
+
     const predictionData = await this.client.request("POST /v1/predictions", {
       version: this.version,
       input,
+      webhook: webhookUrl,
+      webhook_events_filter: webhookEvents,
     });
 
     return new Prediction(predictionData, this);

lib/Model.test.js

@@ -276,4 +276,46 @@ describe("createPrediction()", () => {
       input: { text: "test text" },
     });
   });
+
+  it("supports webhook string", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "testprediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    await model.createPrediction(
+      { text: "test text" },
+      { webhook: "http://test.host/webhook" }
+    );
+
+    expect(client.request).toHaveBeenCalledWith("POST /v1/predictions", {
+      version: "testversion",
+      input: { text: "test text" },
+      webhook: "http://test.host/webhook",
+    });
+  });
+
+  it("supports webhook object", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "testprediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    await model.createPrediction(
+      { text: "test text" },
+      {
+        webhook: {
+          url: "http://test.host/webhook",
+          events: ["output", "completed"],
+        },
+      }
+    );
+
+    expect(client.request).toHaveBeenCalledWith("POST /v1/predictions", {
+      version: "testversion",
+      input: { text: "test text" },
+      webhook: "http://test.host/webhook",
+      webhook_events_filter: ["output", "completed"],
+    });
+  });
 });