Support webhooks when creating predictions

Author: erbridge

README.md

@@ -16,6 +16,8 @@ Install with `npm install replicate`
 
 Set your API token as an environment variable called `REPLICATE_API_TOKEN`.
 
+### Making preedictions
+
 To run a prediction and return its output:
 
 ```js
@@ -75,6 +77,32 @@ 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()`.
 
+#### Webhooks
+
+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",
+    },
+    {
+      // See https://replicate.com/docs/reference/http#create-prediction--webhook
+      webhook: "https://your.host/webhook",
+
+      // See https://replicate.com/docs/reference/http#create-prediction--webhook_events_filter
+      webhookEventsFilter: ["output", "completed"],
+    }
+  );
+```
+
 ## Contributing
 
 While we'd love to accept contributions to this library, please open an issue

lib/Model.js

@@ -122,7 +122,7 @@ export default class Model extends ReplicateObject {
     return prediction;
   }
 
-  async createPrediction(input) {
+  async createPrediction(input, { webhook, webhookEventsFilter } = {}) {
     // 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
@@ -131,6 +131,8 @@ export default class Model extends ReplicateObject {
     const predictionData = await this.client.request("POST /v1/predictions", {
       version: this.version,
       input,
+      webhook,
+      webhook_events_filter: webhookEventsFilter,
     });
 
     return new Prediction(predictionData, this);

lib/Model.test.js

@@ -276,4 +276,44 @@ describe("createPrediction()", () => {
       input: { text: "test text" },
     });
   });
+
+  it("supports webhook URL", 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 events filter", async () => {
+    jest.spyOn(client, "request").mockResolvedValue({
+      id: "testprediction",
+      status: PredictionStatus.SUCCEEDED,
+    });
+
+    await model.createPrediction(
+      { text: "test text" },
+      {
+        webhook: "http://test.host/webhook",
+        webhookEventsFilter: ["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"],
+    });
+  });
 });