Merge branch 'main' into base64-encode-inputs

Author: mattt

.github/workflows/ci.yml

@@ -47,5 +47,8 @@ jobs:
         with:
           node-version: ${{ matrix.node-version }}
           cache: "npm"
-      - run: npm --prefix integration/${{ matrix.suite }} ci --omit=dev
-      - run: npm --prefix integration/${{ matrix.suite }} test
+      # Build a production tarball and run the integration tests against it.
+      - run: |
+          PKG_TARBALL=$(npm --loglevel error pack)
+          npm --prefix integration/${{ matrix.suite }} install "file:/./$PKG_TARBALL"
+          npm --prefix integration/${{ matrix.suite }} test

README.md

@@ -23,8 +23,14 @@ npm install replicate
 Create the client:
 
 ```js
+// CommonJS (default or using .cjs extension)
+const Replicate = require("replicate");
+
+// ESM (where `"module": true` in package.json or using .mjs extension)
 import Replicate from "replicate";
+```
 
+```
 const replicate = new Replicate({
   // get your token from https://replicate.com/account
   auth: "my api token", // defaults to process.env.REPLICATE_API_TOKEN
@@ -69,9 +75,11 @@ console.log(prediction.output);
 
 To run a model that takes a file input, pass a URL to a publicly accessible file. Or, for smaller files (<10MB), you can pass the data directly.
 
-
 ```js
-import { promises as fs } from "fs";
+const fs = require("node:fs/promises");
+
+// Or when using ESM.
+// import fs from "node:fs/promises";
 
 const model = "nightmareai/real-esrgan:42fed1c4974146d4d2414e2be2c5277c7fcf05fcc3a73abf41610695738c1d7b";
 const input = {
@@ -81,6 +89,10 @@ const output = await replicate.run(model, { input });
 // ['https://replicate.delivery/mgxm/e7b0e122-9daa-410e-8cde-006c7308ff4d/output.png']
 ```
 
+## TypeScript
+
+Currently in order to support the module format used by `replicate` you'll need to set `esModuleInterop` to `true` in your tsconfig.json.
+
 ## API
 
 ### Constructor
@@ -112,8 +124,12 @@ you can install a fetch function from an external package like
 and pass it to the `fetch` option in the constructor.
 
 ```js
-import Replicate from "replicate";
-import fetch from "cross-fetch";
+const Replicate = require("replicate");
+const fetch = require("fetch");
+
+// Using ESM:
+// import Replicate from "replicate";
+// import fetch from "cross-fetch";
 
 const replicate = new Replicate({ fetch });
 ```
@@ -188,9 +204,22 @@ Returns `AsyncGenerator<ServerSentEvent>` which yields the events of running the
 Example:
 
 ```js
-for await (const event of replicate.stream("meta/llama-2-70b-chat")) {
-    process.stdout.write(`${event}`);
+const model = "meta/llama-2-70b-chat";
+const options = {
+  input: {
+    prompt: "Write a poem about machine learning in the style of Mary Oliver.",
+  },
+  // webhook: "https://smee.io/dMUlmOMkzeyRGjW" // optional
+};
+const output = [];
+
+for await (const { event, data } of replicate.stream(model, options)) {
+  if (event === "output") {
+    output.push(data);
+  }
 }
+
+console.log(output.join("").trim());
 ```
 
 ### Server-sent events

index.d.ts

@@ -8,13 +8,37 @@ declare module "replicate" {
     response: Response;
   }
 
+  export interface Account {
+    type: "user" | "organization";
+    username: string;
+    name: string;
+    github_url?: string;
+  }
+
   export interface Collection {
     name: string;
     slug: string;
     description: string;
     models?: Model[];
   }
 
+  export interface Deployment {
+    owner: string;
+    name: string;
+    current_release: {
+      number: number;
+      model: string;
+      version: string;
+      created_at: string;
+      created_by: Account;
+      configuration: {
+        hardware: string;
+        min_instances: number;
+        max_instances: number;
+      };
+    };
+  }
+
   export interface Hardware {
     sku: string;
     name: string;
@@ -82,6 +106,10 @@ declare module "replicate" {
     retry?: number;
   }
 
+  export interface WebhookSecret {
+    key: string;
+  }
+
   export default class Replicate {
     constructor(options?: {
       auth?: string;
@@ -140,6 +168,10 @@ declare module "replicate" {
       stop?: (prediction: Prediction) => Promise<boolean>
     ): Promise<Prediction>;
 
+    accounts: {
+      current(): Promise<Account>;
+    };
+
     collections: {
       list(): Promise<Page<Collection>>;
       get(collection_slug: string): Promise<Collection>;
@@ -158,6 +190,10 @@ declare module "replicate" {
           }
         ): Promise<Prediction>;
       };
+      get(
+        deployment_owner: string,
+        deployment_name: string
+      ): Promise<Deployment>;
     };
 
     hardware: {
@@ -222,5 +258,26 @@ declare module "replicate" {
       cancel(training_id: string): Promise<Training>;
       list(): Promise<Page<Training>>;
     };
+
+    webhooks: {
+      default: {
+        secret: {
+          get(): Promise<WebhookSecret>;
+        };
+      };
+    };
   }
+
+  export function validateWebhook(
+    requestData:
+      | Request
+      | {
+          id?: string;
+          timestamp?: string;
+          body: string;
+          secret?: string;
+          signature?: string;
+        },
+    secret: string
+  ): boolean;
 }

index.js

@@ -1,14 +1,16 @@
 const ApiError = require("./lib/error");
 const ModelVersionIdentifier = require("./lib/identifier");
 const { Stream } = require("./lib/stream");
-const { withAutomaticRetries } = require("./lib/util");
+const { withAutomaticRetries, validateWebhook } = require("./lib/util");
 
+const accounts = require("./lib/accounts");
 const collections = require("./lib/collections");
 const deployments = require("./lib/deployments");
 const hardware = require("./lib/hardware");
 const models = require("./lib/models");
 const predictions = require("./lib/predictions");
 const trainings = require("./lib/trainings");
+const webhooks = require("./lib/webhooks");
 
 const packageJSON = require("./package.json");
 
@@ -47,12 +49,17 @@ class Replicate {
     this.baseUrl = options.baseUrl || "https://api.replicate.com/v1";
     this.fetch = options.fetch || globalThis.fetch;
 
+    this.accounts = {
+      current: accounts.current.bind(this),
+    };
+
     this.collections = {
       list: collections.list.bind(this),
       get: collections.get.bind(this),
     };
 
     this.deployments = {
+      get: deployments.get.bind(this),
       predictions: {
         create: deployments.predictions.create.bind(this),
       },
@@ -85,6 +92,14 @@ class Replicate {
       cancel: trainings.cancel.bind(this),
       list: trainings.list.bind(this),
     };
+
+    this.webhooks = {
+      default: {
+        secret: {
+          get: webhooks.default.secret.get.bind(this),
+        },
+      },
+    };
   }
 
   /**
@@ -359,3 +374,4 @@ class Replicate {
 }
 
 module.exports = Replicate;
+module.exports.validateWebhook = validateWebhook;

index.test.ts

@@ -1,5 +1,10 @@
 import { expect, jest, test } from "@jest/globals";
-import Replicate, { ApiError, Model, Prediction } from "replicate";
+import Replicate, {
+  ApiError,
+  Model,
+  Prediction,
+  validateWebhook,
+} from "replicate";
 import nock from "nock";
 import fetch from "cross-fetch";
 
@@ -67,6 +72,22 @@ describe("Replicate client", () => {
     });
   });
 
+  describe("account.get", () => {
+    test("Calls the correct API route", async () => {
+      nock(BASE_URL).get("/account").reply(200, {
+        type: "organization",
+        username: "replicate",
+        name: "Replicate",
+        github_url: "https://github.com/replicate",
+      });
+
+      const account = await client.accounts.current();
+      expect(account.type).toBe("organization");
+      expect(account.username).toBe("replicate");
+    });
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   describe("collections.list", () => {
     test("Calls the correct API route", async () => {
       nock(BASE_URL)
@@ -741,6 +762,47 @@ describe("Replicate client", () => {
     // Add more tests for error handling, edge cases, etc.
   });
 
+  describe("deployments.get", () => {
+    test("Calls the correct API route with the correct payload", async () => {
+      nock(BASE_URL)
+        .get("/deployments/acme/my-app-image-generator")
+        .reply(200, {
+          owner: "acme",
+          name: "my-app-image-generator",
+          current_release: {
+            number: 1,
+            model: "stability-ai/sdxl",
+            version:
+              "da77bc59ee60423279fd632efb4795ab731d9e3ca9705ef3341091fb989b7eaf",
+            created_at: "2024-02-15T16:32:57.018467Z",
+            created_by: {
+              type: "organization",
+              username: "acme",
+              name: "Acme Corp, Inc.",
+              github_url: "https://github.com/acme",
+            },
+            configuration: {
+              hardware: "gpu-t4",
+              scaling: {
+                min_instances: 1,
+                max_instances: 5,
+              },
+            },
+          },
+        });
+
+      const deployment = await client.deployments.get(
+        "acme",
+        "my-app-image-generator"
+      );
+
+      expect(deployment.owner).toBe("acme");
+      expect(deployment.name).toBe("my-app-image-generator");
+      expect(deployment.current_release.model).toBe("stability-ai/sdxl");
+    });
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   describe("predictions.create with model", () => {
     test("Calls the correct API route with the correct payload", async () => {
       nock(BASE_URL)
@@ -1021,5 +1083,39 @@ describe("Replicate client", () => {
     });
   });
 
+  describe("webhooks.default.secret.get", () => {
+    test("Calls the correct API route", async () => {
+      nock(BASE_URL).get("/webhooks/default/secret").reply(200, {
+        key: "whsec_5WbX5kEWLlfzsGNjH64I8lOOqUB6e8FH",
+      });
+
+      const secret = await client.webhooks.default.secret.get();
+      expect(secret.key).toBe("whsec_5WbX5kEWLlfzsGNjH64I8lOOqUB6e8FH");
+    });
+
+    test("Can be used to validate webhook", async () => {
+      // Test case from https://github.com/svix/svix-webhooks/blob/b41728cd98a7e7004a6407a623f43977b82fcba4/javascript/src/webhook.test.ts#L190-L200
+      const request = new Request("http://test.host/webhook", {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          "Webhook-ID": "msg_p5jXN8AQM9LWM0D4loKWxJek",
+          "Webhook-Timestamp": "1614265330",
+          "Webhook-Signature":
+            "v1,g0hM9SsE+OTPJTGt/tmIKtSyZlE3uFJELVlNIOLJ1OE=",
+        },
+        body: `{"test": 2432232314}`,
+      });
+
+      // This is a test secret and should not be used in production
+      const secret = "whsec_MfKQ9r8GKYqrTwjUPD8ILPZIo2LaLaSw";
+
+      const isValid = await validateWebhook(request, secret);
+      expect(isValid).toBe(true);
+    });
+
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   // Continue with tests for other methods
 });

lib/accounts.js

@@ -0,0 +1,16 @@
+/**
+ * Get the current account
+ *
+ * @returns {Promise<object>} Resolves with the current account
+ */
+async function getCurrentAccount() {
+  const response = await this.request("/account", {
+    method: "GET",
+  });
+
+  return response.json();
+}
+
+module.exports = {
+  current: getCurrentAccount,
+};