diff --git a/README.md b/README.md index 1e14ec24..cf12865c 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,13 @@ const requestData = { const webhookIsValid = await validateWebhook(requestData); ``` +> [!NOTE] +> The `validateWebhook` function uses the global `crypto` API available in most JavaScript runtimes. Node <= 18 does not provide this global so in this case you need to either call node with the `--no-experimental-global-webcrypto` or provide the `webcrypto` module manually. +> ```js +> const crypto = require("node:crypto").webcrypto; +> const webhookIsValid = await valdiateWebhook(requestData, crypto); +> ``` + ## TypeScript The `Replicate` constructor and all `replicate.*` methods are fully typed. @@ -343,7 +350,7 @@ const output = await replicate.run(identifier, options, progress); | `identifier` | string | **Required**. The model version identifier in the format `{owner}/{name}:{version}`, for example `stability-ai/sdxl:8beff3369e81422112d93b89ca01426147de542cd4684c244b673b105188fe5f` | | `options.input` | object | **Required**. An object with the model inputs. | | `options.wait` | object | Options for waiting for the prediction to finish | -| `options.wait.type` | `"poll" \| "block"` | `"block"` holds the request open, `"poll"` makes repeated requests to fetch the prediction. Defaults to `"block"` | +| `options.wait.mode` | `"poll" \| "block"` | `"block"` holds the request open, `"poll"` makes repeated requests to fetch the prediction. Defaults to `"block"` | | `options.wait.interval` | number | Polling interval in milliseconds. Defaults to 500 | | `options.wait.timeout` | number | In `"block"` mode determines how long the request will be held open until falling back to polling. Defaults to 60 | | `options.webhook` | string | An HTTPS URL for receiving a webhook when the prediction has new output | @@ -1206,15 +1213,21 @@ Low-level method used by the Replicate client to interact with API endpoints. const response = await replicate.request(route, parameters); ``` -| name | type | description | -| -------------------- | ------ | ------------------------------------------------------------ | -| `options.route` | string | Required. REST API endpoint path. | -| `options.parameters` | object | URL, query, and request body parameters for the given route. | +| name | type | description | +| -------------------- | ------------------- | ----------- | +| `options.route` | `string` | Required. REST API endpoint path. +| `options.params` | `object` | URL query parameters for the given route. | +| `options.method` | `string` | HTTP method for the given route. | +| `options.headers` | `object` | Additional HTTP headers for the given route. | +| `options.data` | `object \| FormData` | Request body. | +| `options.signal` | `AbortSignal` | Optional `AbortSignal`. | The `replicate.request()` method is used by the other methods to interact with the Replicate API. You can call this method directly to make other requests to the API. +The method accepts an `AbortSignal` which can be used to cancel the request in flight. + ### `FileOutput` `FileOutput` is a `ReadableStream` instance that represents a model file output. It can be used to stream file data to disk or as a `Response` body to an HTTP request. diff --git a/biome.json b/biome.json index 094cf0ec..ba1c236b 100644 --- a/biome.json +++ b/biome.json @@ -1,11 +1,7 @@ { "$schema": "https://biomejs.dev/schemas/1.0.0/schema.json", "files": { - "ignore": [ - ".wrangler", - "node_modules", - "vendor/*" - ] + "ignore": [".wrangler", "node_modules", "vendor/*"] }, "formatter": { "indentStyle": "space", diff --git a/index.d.ts b/index.d.ts index 1e417ae7..d4e07840 100644 --- a/index.d.ts +++ b/index.d.ts @@ -1,5 +1,5 @@ declare module "replicate" { - type Status = "starting" | "processing" | "succeeded" | "failed" | "canceled"; + type Status = "starting" | "processing" | "succeeded" | "failed" | "canceled" | "aborted"; type Visibility = "public" | "private"; type WebhookEventType = "start" | "output" | "logs" | "completed"; @@ -19,12 +19,14 @@ declare module "replicate" { username: string; name: string; github_url?: string; + avatar_url?: string; } export interface Collection { name: string; slug: string; description: string; + full_description: string | null; models?: Model[]; } @@ -47,11 +49,11 @@ declare module "replicate" { export interface FileObject { id: string; - name: string; content_type: string; size: number; - etag: string; - checksum: string; + checksums: { + sha256: string; + }; metadata: Record; created_at: string; expires_at: string | null; @@ -71,6 +73,7 @@ declare module "replicate" { name: string; description?: string; visibility: "public" | "private"; + is_official: boolean; github_url?: string; paper_url?: string; license_url?: string; @@ -83,22 +86,26 @@ declare module "replicate" { export interface ModelVersion { id: string; created_at: string; - cog_version: string; - openapi_schema: object; + cog_version: string | null; + openapi_schema: object | null; } export interface Prediction { id: string; status: Status; model: string; - version: string; + version: string | "hidden"; input: object; output?: any; // TODO: this should be `unknown` source: "api" | "web"; error?: unknown; logs?: string; + data_removed: boolean; + deadline?: string; + deployment?: string; metrics?: { predict_time?: number; + total_time?: number; }; webhook?: string; webhook_events_filter?: WebhookEventType[]; @@ -109,10 +116,38 @@ declare module "replicate" { get: string; cancel: string; stream?: string; + web?: string; }; } - export type Training = Prediction; + export interface Training { + id: string; + status: Status; + model: string; + version: string; + input: object; + output?: { + version?: string; + weights?: string; + }; + source: "api" | "web"; + error?: unknown; + logs?: string; + metrics?: { + predict_time?: number; + total_time?: number; + }; + webhook?: string; + webhook_events_filter?: WebhookEventType[]; + created_at: string; + started_at?: string; + completed_at?: string; + urls: { + get: string; + cancel: string; + web?: string; + }; + } export type FileEncodingStrategy = "default" | "upload" | "data-uri"; @@ -173,6 +208,7 @@ declare module "replicate" { webhook?: string; webhook_events_filter?: WebhookEventType[]; signal?: AbortSignal; + useFileOutput?: boolean; } ): AsyncGenerator; @@ -183,10 +219,14 @@ declare module "replicate" { headers?: object | Headers; params?: object; data?: object; + signal?: AbortSignal; } ): Promise; - paginate(endpoint: () => Promise>): AsyncGenerator<[T]>; + paginate( + endpoint: () => Promise>, + options?: { signal?: AbortSignal } + ): AsyncGenerator; wait( prediction: Prediction, @@ -197,12 +237,15 @@ declare module "replicate" { ): Promise; accounts: { - current(): Promise; + current(options?: { signal?: AbortSignal }): Promise; }; collections: { - list(): Promise>; - get(collection_slug: string): Promise; + list(options?: { signal?: AbortSignal }): Promise>; + get( + collection_slug: string, + options?: { signal?: AbortSignal } + ): Promise; }; deployments: { @@ -217,21 +260,26 @@ declare module "replicate" { webhook?: string; webhook_events_filter?: WebhookEventType[]; wait?: number | boolean; + signal?: AbortSignal; } ): Promise; }; get( deployment_owner: string, - deployment_name: string + deployment_name: string, + options?: { signal?: AbortSignal } + ): Promise; + create( + deployment_config: { + name: string; + model: string; + version: string; + hardware: string; + min_instances: number; + max_instances: number; + }, + options?: { signal?: AbortSignal } ): Promise; - create(deployment_config: { - name: string; - model: string; - version: string; - hardware: string; - min_instances: number; - max_instances: number; - }): Promise; update( deployment_owner: string, deployment_name: string, @@ -245,32 +293,45 @@ declare module "replicate" { | { hardware: string } | { min_instances: number } | { max_instances: number } - ) + ), + options?: { signal?: AbortSignal } ): Promise; delete( deployment_owner: string, - deployment_name: string + deployment_name: string, + options?: { signal?: AbortSignal } ): Promise; - list(): Promise>; + list(options?: { signal?: AbortSignal }): Promise>; }; files: { create( file: Blob | Buffer, - metadata?: Record + metadata?: Record, + options?: { signal?: AbortSignal } + ): Promise; + list(options?: { signal?: AbortSignal }): Promise>; + get( + file_id: string, + options?: { signal?: AbortSignal } ): Promise; - list(): Promise>; - get(file_id: string): Promise; - delete(file_id: string): Promise; + delete( + file_id: string, + options?: { signal?: AbortSignal } + ): Promise; }; hardware: { - list(): Promise; + list(options?: { signal?: AbortSignal }): Promise; }; models: { - get(model_owner: string, model_name: string): Promise; - list(): Promise>; + get( + model_owner: string, + model_name: string, + options?: { signal?: AbortSignal } + ): Promise; + list(options?: { signal?: AbortSignal }): Promise>; create( model_owner: string, model_name: string, @@ -282,17 +343,26 @@ declare module "replicate" { paper_url?: string; license_url?: string; cover_image_url?: string; + signal?: AbortSignal; } ): Promise; versions: { - list(model_owner: string, model_name: string): Promise; + list( + model_owner: string, + model_name: string, + options?: { signal?: AbortSignal } + ): Promise>; get( model_owner: string, model_name: string, - version_id: string + version_id: string, + options?: { signal?: AbortSignal } ): Promise; }; - search(query: string): Promise>; + search( + query: string, + options?: { signal?: AbortSignal } + ): Promise>; }; predictions: { @@ -306,11 +376,18 @@ declare module "replicate" { webhook?: string; webhook_events_filter?: WebhookEventType[]; wait?: boolean | number; + signal?: AbortSignal; } & ({ version: string } | { model: string }) ): Promise; - get(prediction_id: string): Promise; - cancel(prediction_id: string): Promise; - list(): Promise>; + get( + prediction_id: string, + options?: { signal?: AbortSignal } + ): Promise; + cancel( + prediction_id: string, + options?: { signal?: AbortSignal } + ): Promise; + list(options?: { signal?: AbortSignal }): Promise>; }; trainings: { @@ -323,33 +400,44 @@ declare module "replicate" { input: object; webhook?: string; webhook_events_filter?: WebhookEventType[]; + signal?: AbortSignal; } ): Promise; - get(training_id: string): Promise; - cancel(training_id: string): Promise; - list(): Promise>; + get( + training_id: string, + options?: { signal?: AbortSignal } + ): Promise; + cancel( + training_id: string, + options?: { signal?: AbortSignal } + ): Promise; + list(options?: { signal?: AbortSignal }): Promise>; }; webhooks: { default: { secret: { - get(): Promise; + get(options?: { signal?: AbortSignal }): Promise; }; }; }; } export function validateWebhook( - requestData: - | Request - | { - id?: string; - timestamp?: string; - body: string; - secret?: string; - signature?: string; - }, - secret: string + request: Request, + secret: string, + crypto?: Crypto + ): Promise; + + export function validateWebhook( + requestData: { + id: string; + timestamp: string; + signature: string; + body: string; + secret: string; + }, + crypto?: Crypto ): Promise; export function parseProgressFromLogs(logs: Prediction | string): { diff --git a/index.js b/index.js index a5755d91..3fc14dd5 100644 --- a/index.js +++ b/index.js @@ -225,6 +225,7 @@ class Replicate { * @param {object} [options.params] - Query parameters * @param {object|Headers} [options.headers] - HTTP headers * @param {object} [options.data] - Body parameters + * @param {AbortSignal} [options.signal] - AbortSignal to cancel the request * @returns {Promise} - Resolves with the response object * @throws {ApiError} If the request failed */ @@ -241,7 +242,7 @@ class Replicate { ); } - const { method = "GET", params = {}, data } = options; + const { method = "GET", params = {}, data, signal } = options; for (const [key, value] of Object.entries(params)) { url.searchParams.append(key, value); @@ -273,6 +274,7 @@ class Replicate { method, headers, body, + signal, }; const shouldRetry = @@ -313,7 +315,12 @@ class Replicate { * @yields {ServerSentEvent} Each streamed event from the prediction */ async *stream(ref, options) { - const { wait, signal, ...data } = options; + const { + wait, + signal, + useFileOutput = this.useFileOutput, + ...data + } = options; const identifier = ModelVersionIdentifier.parse(ref); @@ -336,7 +343,10 @@ class Replicate { const stream = createReadableStream({ url: prediction.urls.stream, fetch: this.fetch, - ...(signal ? { options: { signal } } : {}), + options: { + useFileOutput, + ...(signal ? { signal } : {}), + }, }); yield* streamAsyncIterator(stream); @@ -354,15 +364,20 @@ class Replicate { * console.log(page); * } * @param {Function} endpoint - Function that returns a promise for the next page of results + * @param {object} [options] + * @param {AbortSignal} [options.signal] - AbortSignal to cancel the request. * @yields {object[]} Each page of results */ - async *paginate(endpoint) { + async *paginate(endpoint, options = {}) { const response = await endpoint(); yield response.results; - if (response.next) { + if (response.next && !(options.signal && options.signal.aborted)) { const nextPage = () => - this.request(response.next, { method: "GET" }).then((r) => r.json()); - yield* this.paginate(nextPage); + this.request(response.next, { + method: "GET", + signal: options.signal, + }).then((r) => r.json()); + yield* this.paginate(nextPage, options); } } diff --git a/index.test.ts b/index.test.ts index c67035a3..96f50db7 100644 --- a/index.test.ts +++ b/index.test.ts @@ -4,11 +4,13 @@ import Replicate, { FileOutput, Model, Prediction, + Training, validateWebhook, parseProgressFromLogs, } from "replicate"; import nock from "nock"; import { createReadableStream } from "./lib/stream"; +import { webcrypto } from "node:crypto"; let client: Replicate; const BASE_URL = "https://api.replicate.com/v1"; @@ -98,6 +100,90 @@ describe("Replicate client", () => { }); }); + describe("paginate", () => { + test("pages through results", async () => { + nock(BASE_URL) + .get("/collections") + .reply(200, { + results: [ + { + name: "Super resolution", + slug: "super-resolution", + description: + "Upscaling models that create high-quality images from low-quality images.", + }, + ], + next: `${BASE_URL}/collections?page=2`, + previous: null, + }); + nock(BASE_URL) + .get("/collections?page=2") + .reply(200, { + results: [ + { + name: "Image classification", + slug: "image-classification", + description: "Models that classify images.", + }, + ], + next: null, + previous: null, + }); + + const iterator = client.paginate(client.collections.list); + + const firstPage = (await iterator.next()).value; + expect(firstPage.length).toBe(1); + + const secondPage = (await iterator.next()).value; + expect(secondPage.length).toBe(1); + }); + + test("accepts an abort signal", async () => { + nock(BASE_URL) + .get("/collections") + .reply(200, { + results: [ + { + name: "Super resolution", + slug: "super-resolution", + description: + "Upscaling models that create high-quality images from low-quality images.", + }, + ], + next: `${BASE_URL}/collections?page=2`, + previous: null, + }); + nock(BASE_URL) + .get("/collections?page=2") + .reply(200, { + results: [ + { + name: "Image classification", + slug: "image-classification", + description: "Models that classify images.", + }, + ], + next: null, + previous: null, + }); + + const controller = new AbortController(); + const iterator = client.paginate(client.collections.list, { + signal: controller.signal, + }); + + const firstIteration = await iterator.next(); + expect(firstIteration.value.length).toBe(1); + + controller.abort(); + + const secondIteration = await iterator.next(); + expect(secondIteration.value).toBeUndefined(); + expect(secondIteration.done).toBe(true); + }); + }); + describe("account.get", () => { test("Calls the correct API route", async () => { nock(BASE_URL).get("/account").reply(200, { @@ -821,7 +907,7 @@ describe("Replicate client", () => { next: null, }); - const results: Prediction[] = []; + const results: Training[] = []; for await (const batch of client.paginate(client.trainings.list)) { results.push(...batch); } @@ -1091,11 +1177,11 @@ describe("Replicate client", () => { .post("/files") .reply(200, { id: "123", - name: "test-file", content_type: "application/octet-stream", size: 1024, - etag: "abc123", - checksum: "sha256:1234567890abcdef", + checksums: { + sha256: "1234567890abcdef", + }, metadata: {}, created_at: "2023-01-01T00:00:00Z", expires_at: null, @@ -1105,7 +1191,6 @@ describe("Replicate client", () => { }); const file = await client.files.create(testCase.value); expect(file.id).toBe("123"); - expect(file.name).toBe("test-file"); } }); }); @@ -1116,11 +1201,11 @@ describe("Replicate client", () => { .get("/files/123") .reply(200, { id: "123", - name: "test-file", content_type: "application/octet-stream", size: 1024, - etag: "abc123", - checksum: "sha256:1234567890abcdef", + checksums: { + sha256: "1234567890abcdef", + }, metadata: {}, created_at: "2023-01-01T00:00:00Z", expires_at: null, @@ -1131,7 +1216,6 @@ describe("Replicate client", () => { const file = await client.files.get("123"); expect(file.id).toBe("123"); - expect(file.name).toBe("test-file"); }); }); @@ -1145,11 +1229,11 @@ describe("Replicate client", () => { results: [ { id: "123", - name: "test-file", content_type: "application/octet-stream", size: 1024, - etag: "abc123", - checksum: "sha256:1234567890abcdef", + checksums: { + sha256: "1234567890abcdef", + }, metadata: {}, created_at: "2023-01-01T00:00:00Z", expires_at: null, @@ -1779,9 +1863,40 @@ describe("Replicate client", () => { // This is a test secret and should not be used in production const secret = "whsec_MfKQ9r8GKYqrTwjUPD8ILPZIo2LaLaSw"; + if (globalThis.crypto) { + const isValid = await validateWebhook(request, secret); + expect(isValid).toBe(true); + } else { + const isValid = await validateWebhook( + request, + secret, + webcrypto as Crypto // Node 18 requires this to be passed manually + ); + expect(isValid).toBe(true); + } + }); + + 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 requestData = { + id: "msg_p5jXN8AQM9LWM0D4loKWxJek", + timestamp: "1614265330", + signature: "v1,g0hM9SsE+OTPJTGt/tmIKtSyZlE3uFJELVlNIOLJ1OE=", + body: `{"test": 2432232314}`, + secret: "whsec_MfKQ9r8GKYqrTwjUPD8ILPZIo2LaLaSw", + }; - const isValid = await validateWebhook(request, secret); - expect(isValid).toBe(true); + // This is a test secret and should not be used in production + if (globalThis.crypto) { + const isValid = await validateWebhook(requestData); + expect(isValid).toBe(true); + } else { + const isValid = await validateWebhook( + requestData, + webcrypto as Crypto // Node 18 requires this to be passed manually + ); + expect(isValid).toBe(true); + } }); // Add more tests for error handling, edge cases, etc. @@ -1790,8 +1905,12 @@ describe("Replicate client", () => { // Continue with tests for other methods describe("createReadableStream", () => { - function createStream(body: string | ReadableStream, status = 200) { - const streamEndpoint = "https://stream.replicate.com/fake_stream"; + function createStream( + body: string | ReadableStream, + status = 200, + streamEndpoint = "https://stream.replicate.com/fake_stream", + options: { useFileOutput?: boolean } = {} + ) { const fetch = jest.fn((url) => { if (url !== streamEndpoint) { throw new Error(`Unmocked call to fetch() with url: ${url}`); @@ -1801,6 +1920,7 @@ describe("Replicate client", () => { return createReadableStream({ url: streamEndpoint, fetch: fetch as any, + options, }); } @@ -2077,5 +2197,95 @@ describe("Replicate client", () => { ); expect(await iterator.next()).toEqual({ done: true }); }); + + describe("file streams", () => { + test("emits FileOutput objects", async () => { + const stream = createStream( + ` + event: output + id: EVENT_1 + data: data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + + event: output + id: EVENT_2 + data: https://delivery.replicate.com/my_file.png + + event: done + id: EVENT_3 + data: {} + + `.replace(/^[ ]+/gm, ""), + 200, + "https://stream.replicate.com/v1/files/abcd" + ); + + const iterator = stream[Symbol.asyncIterator](); + const { value: event1 } = await iterator.next(); + expect(event1.data).toBeInstanceOf(ReadableStream); + expect(event1.data.url().href).toEqual( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ); + + const { value: event2 } = await iterator.next(); + expect(event2.data).toBeInstanceOf(ReadableStream); + expect(event2.data.url().href).toEqual( + "https://delivery.replicate.com/my_file.png" + ); + + expect(await iterator.next()).toEqual({ + done: false, + value: { event: "done", id: "EVENT_3", data: "{}" }, + }); + + expect(await iterator.next()).toEqual({ done: true }); + }); + + test("emits strings when useFileOutput is false", async () => { + const stream = createStream( + ` + event: output + id: EVENT_1 + data: data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + + event: output + id: EVENT_2 + data: https://delivery.replicate.com/my_file.png + + event: done + id: EVENT_3 + data: {} + + `.replace(/^[ ]+/gm, ""), + 200, + "https://stream.replicate.com/v1/files/abcd", + { useFileOutput: false } + ); + + const iterator = stream[Symbol.asyncIterator](); + + expect(await iterator.next()).toEqual({ + done: false, + value: { + event: "output", + id: "EVENT_1", + data: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==", + }, + }); + expect(await iterator.next()).toEqual({ + done: false, + value: { + event: "output", + id: "EVENT_2", + data: "https://delivery.replicate.com/my_file.png", + }, + }); + expect(await iterator.next()).toEqual({ + done: false, + value: { event: "done", id: "EVENT_3", data: "{}" }, + }); + + expect(await iterator.next()).toEqual({ done: true }); + }); + }); }); }); diff --git a/integration/next/pages/index.js b/integration/next/pages/index.js index fc9581a7..09124388 100644 --- a/integration/next/pages/index.js +++ b/integration/next/pages/index.js @@ -1,5 +1,5 @@ export default () => ( -
-

Welcome to Next.js

-
-) +
+

Welcome to Next.js

+
+); diff --git a/lib/accounts.js b/lib/accounts.js index b3bbd9fe..72a94afa 100644 --- a/lib/accounts.js +++ b/lib/accounts.js @@ -1,11 +1,14 @@ /** * Get the current account * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the current account */ -async function getCurrentAccount() { +async function getCurrentAccount({ signal } = {}) { const response = await this.request("/account", { method: "GET", + signal, }); return response.json(); diff --git a/lib/collections.js b/lib/collections.js index 9332aaab..7b8e8f11 100644 --- a/lib/collections.js +++ b/lib/collections.js @@ -2,11 +2,14 @@ * Fetch a model collection * * @param {string} collection_slug - Required. The slug of the collection. See http://replicate.com/collections + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with the collection data */ -async function getCollection(collection_slug) { +async function getCollection(collection_slug, { signal } = {}) { const response = await this.request(`/collections/${collection_slug}`, { method: "GET", + signal, }); return response.json(); @@ -15,11 +18,14 @@ async function getCollection(collection_slug) { /** * Fetch a list of model collections * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with the collections data */ -async function listCollections() { +async function listCollections({ signal } = {}) { const response = await this.request("/collections", { method: "GET", + signal, }); return response.json(); diff --git a/lib/deployments.js b/lib/deployments.js index 716c8e19..d45c4f39 100644 --- a/lib/deployments.js +++ b/lib/deployments.js @@ -10,10 +10,11 @@ const { transformFileInputs } = require("./util"); * @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|integer} [options.wait] - Whether to wait until the prediction is completed before returning. If an integer is provided, it will wait for that many seconds. Defaults to false + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the created prediction data */ async function createPrediction(deployment_owner, deployment_name, options) { - const { input, wait, ...data } = options; + const { input, wait, signal, ...data } = options; if (data.webhook) { try { @@ -47,6 +48,7 @@ async function createPrediction(deployment_owner, deployment_name, options) { this.fileEncodingStrategy ), }, + signal, } ); @@ -58,13 +60,20 @@ async function createPrediction(deployment_owner, deployment_name, options) { * * @param {string} deployment_owner - Required. The username of the user or organization who owns the deployment * @param {string} deployment_name - Required. The name of the deployment + * @param {object] [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the deployment data */ -async function getDeployment(deployment_owner, deployment_name) { +async function getDeployment( + deployment_owner, + deployment_name, + { signal } = {} +) { const response = await this.request( `/deployments/${deployment_owner}/${deployment_name}`, { method: "GET", + signal, } ); @@ -84,13 +93,16 @@ async function getDeployment(deployment_owner, deployment_name) { /** * Create a deployment * - * @param {DeploymentCreateRequest} config - Required. The deployment config. + * @param {DeploymentCreateRequest} deployment_config - Required. The deployment config. + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the deployment data */ -async function createDeployment(deployment_config) { +async function createDeployment(deployment_config, { signal } = {}) { const response = await this.request("/deployments", { method: "POST", data: deployment_config, + signal, }); return response.json(); @@ -110,18 +122,22 @@ async function createDeployment(deployment_config) { * @param {string} deployment_owner - Required. The username of the user or organization who owns the deployment * @param {string} deployment_name - Required. The name of the deployment * @param {DeploymentUpdateRequest} deployment_config - Required. The deployment changes. + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the deployment data */ async function updateDeployment( deployment_owner, deployment_name, - deployment_config + deployment_config, + { signal } = {} ) { const response = await this.request( `/deployments/${deployment_owner}/${deployment_name}`, { method: "PATCH", data: deployment_config, + signal, } ); @@ -133,13 +149,20 @@ async function updateDeployment( * * @param {string} deployment_owner - Required. The username of the user or organization who owns the deployment * @param {string} deployment_name - Required. The name of the deployment + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with true if the deployment was deleted */ -async function deleteDeployment(deployment_owner, deployment_name) { +async function deleteDeployment( + deployment_owner, + deployment_name, + { signal } = {} +) { const response = await this.request( `/deployments/${deployment_owner}/${deployment_name}`, { method: "DELETE", + signal, } ); @@ -149,11 +172,14 @@ async function deleteDeployment(deployment_owner, deployment_name) { /** * List all deployments * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with a page of deployments */ -async function listDeployments() { +async function listDeployments({ signal } = {}) { const response = await this.request("/deployments", { method: "GET", + signal, }); return response.json(); diff --git a/lib/files.js b/lib/files.js index de49c580..c8101398 100644 --- a/lib/files.js +++ b/lib/files.js @@ -3,9 +3,11 @@ * * @param {object} file - Required. The file object. * @param {object} metadata - Optional. User-provided metadata associated with the file. + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with the file data */ -async function createFile(file, metadata = {}) { +async function createFile(file, metadata = {}, { signal } = {}) { const form = new FormData(); let filename; @@ -36,6 +38,7 @@ async function createFile(file, metadata = {}) { headers: { "Content-Type": "multipart/form-data", }, + signal, }); return response.json(); @@ -44,11 +47,14 @@ async function createFile(file, metadata = {}) { /** * List all files * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with the files data */ -async function listFiles() { +async function listFiles({ signal } = {}) { const response = await this.request("/files", { method: "GET", + signal, }); return response.json(); @@ -58,11 +64,14 @@ async function listFiles() { * Get a file * * @param {string} file_id - Required. The ID of the file. + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with the file data */ -async function getFile(file_id) { +async function getFile(file_id, { signal } = {}) { const response = await this.request(`/files/${file_id}`, { method: "GET", + signal, }); return response.json(); @@ -72,11 +81,14 @@ async function getFile(file_id) { * Delete a file * * @param {string} file_id - Required. The ID of the file. + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with true if the file was deleted */ -async function deleteFile(file_id) { +async function deleteFile(file_id, { signal } = {}) { const response = await this.request(`/files/${file_id}`, { method: "DELETE", + signal, }); return response.status === 204; diff --git a/lib/hardware.js b/lib/hardware.js index d717548f..e981b1fa 100644 --- a/lib/hardware.js +++ b/lib/hardware.js @@ -1,11 +1,14 @@ /** * List hardware * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the array of hardware */ -async function listHardware() { +async function listHardware({ signal } = {}) { const response = await this.request("/hardware", { method: "GET", + signal, }); return response.json(); diff --git a/lib/models.js b/lib/models.js index 272d9edb..4a3fcdde 100644 --- a/lib/models.js +++ b/lib/models.js @@ -3,11 +3,14 @@ * * @param {string} model_owner - Required. The name of the user or organization that owns the model * @param {string} model_name - Required. The name of the model + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the model data */ -async function getModel(model_owner, model_name) { +async function getModel(model_owner, model_name, { signal } = {}) { const response = await this.request(`/models/${model_owner}/${model_name}`, { method: "GET", + signal, }); return response.json(); @@ -18,13 +21,16 @@ async function getModel(model_owner, model_name) { * * @param {string} model_owner - Required. The name of the user or organization that owns the model * @param {string} model_name - Required. The name of the model + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the list of model versions */ -async function listModelVersions(model_owner, model_name) { +async function listModelVersions(model_owner, model_name, { signal } = {}) { const response = await this.request( `/models/${model_owner}/${model_name}/versions`, { method: "GET", + signal, } ); @@ -37,13 +43,21 @@ async function listModelVersions(model_owner, model_name) { * @param {string} model_owner - Required. The name of the user or organization that owns the model * @param {string} model_name - Required. The name of the model * @param {string} version_id - Required. The model version + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the model version data */ -async function getModelVersion(model_owner, model_name, version_id) { +async function getModelVersion( + model_owner, + model_name, + version_id, + { signal } = {} +) { const response = await this.request( `/models/${model_owner}/${model_name}/versions/${version_id}`, { method: "GET", + signal, } ); @@ -53,11 +67,14 @@ async function getModelVersion(model_owner, model_name, version_id) { /** * List all public models * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the model version data */ -async function listModels() { +async function listModels({ signal } = {}) { const response = await this.request("/models", { method: "GET", + signal, }); return response.json(); @@ -76,14 +93,17 @@ async function listModels() { * @param {string} options.paper_url - A URL for the model's paper. * @param {string} options.license_url - A URL for the model's license. * @param {string} options.cover_image_url - A URL for the model's cover image. This should be an image file. + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the model version data */ async function createModel(model_owner, model_name, options) { - const data = { owner: model_owner, name: model_name, ...options }; + const { signal, ...rest } = options; + const data = { owner: model_owner, name: model_name, ...rest }; const response = await this.request("/models", { method: "POST", data, + signal, }); return response.json(); @@ -93,15 +113,18 @@ async function createModel(model_owner, model_name, options) { * Search for public models * * @param {string} query - The search query + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with a page of models matching the search query */ -async function search(query) { +async function search(query, { signal } = {}) { const response = await this.request("/models", { method: "QUERY", headers: { "Content-Type": "text/plain", }, data: query, + signal, }); return response.json(); diff --git a/lib/predictions.js b/lib/predictions.js index f8e1c5ac..708d04b3 100644 --- a/lib/predictions.js +++ b/lib/predictions.js @@ -10,10 +10,11 @@ const { transformFileInputs } = require("./util"); * @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|integer} [options.wait] - Whether to wait until the prediction is completed before returning. If an integer is provided, it will wait for that many seconds. Defaults to false + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the created prediction */ async function createPrediction(options) { - const { model, version, input, wait, ...data } = options; + const { model, version, input, wait, signal, ...data } = options; if (data.webhook) { try { @@ -48,6 +49,7 @@ async function createPrediction(options) { ), version, }, + signal, }); } else if (model) { response = await this.request(`/models/${model}/predictions`, { @@ -61,6 +63,7 @@ async function createPrediction(options) { this.fileEncodingStrategy ), }, + signal, }); } else { throw new Error("Either model or version must be specified"); @@ -73,11 +76,14 @@ async function createPrediction(options) { * Fetch a prediction by ID * * @param {number} prediction_id - Required. The prediction ID + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the prediction data */ -async function getPrediction(prediction_id) { +async function getPrediction(prediction_id, { signal } = {}) { const response = await this.request(`/predictions/${prediction_id}`, { method: "GET", + signal, }); return response.json(); @@ -87,11 +93,14 @@ async function getPrediction(prediction_id) { * Cancel a prediction by ID * * @param {string} prediction_id - Required. The training ID + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the data for the training */ -async function cancelPrediction(prediction_id) { +async function cancelPrediction(prediction_id, { signal } = {}) { const response = await this.request(`/predictions/${prediction_id}/cancel`, { method: "POST", + signal, }); return response.json(); @@ -100,11 +109,14 @@ async function cancelPrediction(prediction_id) { /** * List all predictions * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with a page of predictions */ -async function listPredictions() { +async function listPredictions({ signal } = {}) { const response = await this.request("/predictions", { method: "GET", + signal, }); return response.json(); diff --git a/lib/stream.js b/lib/stream.js index 2c899bd9..802a98e1 100644 --- a/lib/stream.js +++ b/lib/stream.js @@ -53,6 +53,7 @@ class ServerSentEvent { */ function createReadableStream({ url, fetch, options = {} }) { const { useFileOutput = true, headers = {}, ...initOptions } = options; + const shouldProcessFileOutput = useFileOutput && isFileStream(url); return new ReadableStream({ async start(controller) { @@ -89,11 +90,11 @@ function createReadableStream({ url, fetch, options = {} }) { let data = event.data; if ( - useFileOutput && - typeof data === "string" && - (data.startsWith("https:") || data.startsWith("data:")) + event.event === "output" && + shouldProcessFileOutput && + typeof data === "string" ) { - data = createFileOutput({ data, fetch }); + data = createFileOutput({ url: data, fetch }); } controller.enqueue(new ServerSentEvent(event.event, data, event.id)); @@ -169,6 +170,13 @@ function createFileOutput({ url, fetch }) { }); } +function isFileStream(url) { + try { + return new URL(url).pathname.startsWith("/v1/files/"); + } catch {} + return false; +} + module.exports = { createFileOutput, createReadableStream, diff --git a/lib/trainings.js b/lib/trainings.js index 6b13dca2..49640b9e 100644 --- a/lib/trainings.js +++ b/lib/trainings.js @@ -9,10 +9,11 @@ * @param {object} options.input - Required. An object with the model inputs * @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the training updates * @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`) + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the data for the created training */ async function createTraining(model_owner, model_name, version_id, options) { - const { ...data } = options; + const { signal, ...data } = options; if (data.webhook) { try { @@ -28,6 +29,7 @@ async function createTraining(model_owner, model_name, version_id, options) { { method: "POST", data, + signal, } ); @@ -38,11 +40,14 @@ async function createTraining(model_owner, model_name, version_id, options) { * Fetch a training by ID * * @param {string} training_id - Required. The training ID + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the data for the training */ -async function getTraining(training_id) { +async function getTraining(training_id, { signal } = {}) { const response = await this.request(`/trainings/${training_id}`, { method: "GET", + signal, }); return response.json(); @@ -52,11 +57,14 @@ async function getTraining(training_id) { * Cancel a training by ID * * @param {string} training_id - Required. The training ID + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the data for the training */ -async function cancelTraining(training_id) { +async function cancelTraining(training_id, { signal } = {}) { const response = await this.request(`/trainings/${training_id}/cancel`, { method: "POST", + signal, }); return response.json(); @@ -65,11 +73,14 @@ async function cancelTraining(training_id) { /** * List all trainings * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} - Resolves with a page of trainings */ -async function listTrainings() { +async function listTrainings({ signal } = {}) { const response = await this.request("/trainings", { method: "GET", + signal, }); return response.json(); diff --git a/lib/util.js b/lib/util.js index bd3c31e9..09665774 100644 --- a/lib/util.js +++ b/lib/util.js @@ -10,6 +10,7 @@ const { create: createFile } = require("./files"); * @param {string} requestData.body - The raw body of the incoming webhook request. * @param {string} requestData.secret - The webhook secret, obtained from `replicate.webhooks.defaul.secret` method. * @param {string} requestData.signature - The webhook signature header from the incoming request, comprising one or more space-delimited signatures. + * @param {Crypto} [crypto] - An optional `Crypto` implementation that conforms to the [browser Crypto interface](https://developer.mozilla.org/en-US/docs/Web/API/Window/crypto) */ /** @@ -22,6 +23,7 @@ const { create: createFile } = require("./files"); * @param {string} requestData.headers["webhook-signature"] - The webhook signature header from the incoming request, comprising one or more space-delimited signatures * @param {string} requestData.body - The raw body of the incoming webhook request * @param {string} secret - The webhook secret, obtained from `replicate.webhooks.defaul.secret` method + * @param {Crypto} [crypto] - An optional `Crypto` implementation that conforms to the [browser Crypto interface](https://developer.mozilla.org/en-US/docs/Web/API/Window/crypto) */ /** @@ -30,9 +32,13 @@ const { create: createFile } = require("./files"); * @returns {Promise} - True if the signature is valid * @throws {Error} - If the request is missing required headers, body, or secret */ -async function validateWebhook(requestData, secret) { - let { id, timestamp, body, signature } = requestData; - const signingSecret = secret || requestData.secret; +async function validateWebhook(requestData, secretOrCrypto, customCrypto) { + let id; + let body; + let timestamp; + let signature; + let secret; + let crypto = globalThis.crypto; if (requestData && requestData.headers && requestData.body) { if (typeof requestData.headers.get === "function") { @@ -47,6 +53,25 @@ async function validateWebhook(requestData, secret) { signature = requestData.headers["webhook-signature"]; } body = requestData.body; + if (typeof secretOrCrypto !== "string") { + throw new Error( + "Unexpected value for secret passed to validateWebhook, expected a string" + ); + } + + secret = secretOrCrypto; + if (customCrypto) { + crypto = customCrypto; + } + } else { + id = requestData.id; + body = requestData.body; + timestamp = requestData.timestamp; + signature = requestData.signature; + secret = requestData.secret; + if (secretOrCrypto) { + crypto = secretOrCrypto; + } } if (body instanceof ReadableStream || body.readable) { @@ -71,15 +96,22 @@ async function validateWebhook(requestData, secret) { throw new Error("Missing required body"); } - if (!signingSecret) { + if (!secret) { throw new Error("Missing required secret"); } + if (!crypto) { + throw new Error( + 'Missing `crypto` implementation. If using Node 18 pass in require("node:crypto").webcrypto' + ); + } + const signedContent = `${id}.${timestamp}.${body}`; const computedSignature = await createHMACSHA256( - signingSecret.split("_").pop(), - signedContent + secret.split("_").pop(), + signedContent, + crypto ); const expectedSignatures = signature @@ -94,27 +126,10 @@ async function validateWebhook(requestData, secret) { /** * @param {string} secret - base64 encoded string * @param {string} data - text body of request + * @param {Crypto} crypto - an implementation of the web Crypto api */ -async function createHMACSHA256(secret, data) { +async function createHMACSHA256(secret, data, crypto) { const encoder = new TextEncoder(); - let crypto = globalThis.crypto; - - // In Node 18 the `crypto` global is behind a --no-experimental-global-webcrypto flag - if (typeof crypto === "undefined" && typeof require === "function") { - // NOTE: Webpack (primarily as it's used by Next.js) and perhaps some - // other bundlers do not currently support the `node` protocol and will - // error if it's found in the source. Other platforms like CloudFlare - // will only support requires when using the node protocol. - // - // As this line is purely to support Node 18.x we make an indirect request - // to the require function which fools Webpack... - // - // We may be able to remove this in future as it looks like Webpack is getting - // support for requiring using the `node:` protocol. - // See: https://github.com/webpack/webpack/issues/18277 - crypto = require.call(null, "node:crypto").webcrypto; - } - const key = await crypto.subtle.importKey( "raw", base64ToBytes(secret), @@ -232,7 +247,7 @@ async function withAutomaticRetries(request, options = {}) { * @param {Replicate} client - The client used to upload the file * @param {object} inputs - The inputs to transform * @param {"default" | "upload" | "data-uri"} strategy - Whether to upload files to Replicate, encode as dataURIs or try both. - * @returns {object} - The transformed inputs + * @returns {Promise} - The transformed inputs * @throws {ApiError} If the request to upload the file fails */ async function transformFileInputs(client, inputs, strategy) { @@ -265,7 +280,7 @@ async function transformFileInputs(client, inputs, strategy) { * * @param {Replicate} client - The client used to upload the file * @param {object} inputs - The inputs to transform - * @returns {object} - The transformed inputs + * @returns {Promise} - The transformed inputs * @throws {ApiError} If the request to upload the file fails */ async function transformFileInputsToReplicateFileURLs(client, inputs) { @@ -286,8 +301,8 @@ const MAX_DATA_URI_SIZE = 10_000_000; * base64-encoded data URI. * * @param {object} inputs - The inputs to transform - * @returns {object} - The transformed inputs - * @throws {Error} If the size of inputs exceeds a given threshould set by MAX_DATA_URI_SIZE + * @returns {Promise} - The transformed inputs + * @throws {Error} If the size of inputs exceeds a given threshold set by MAX_DATA_URI_SIZE */ async function transformFileInputsToBase64EncodedDataURIs(inputs) { let totalBytes = 0; @@ -296,10 +311,10 @@ async function transformFileInputsToBase64EncodedDataURIs(inputs) { let mime; if (value instanceof Blob) { - // Currently we use a NodeJS only API for base64 encoding, as + // Currently, we use a NodeJS only API for base64 encoding, as // we move to support the browser we could support either using // btoa (which does string encoding), the FileReader API or - // a JavaScript implenentation like base64-js. + // a JavaScript implementation like base64-js. // See: https://developer.mozilla.org/en-US/docs/Glossary/Base64 // See: https://github.com/beatgammit/base64-js buffer = await value.arrayBuffer(); diff --git a/lib/webhooks.js b/lib/webhooks.js index f1324ec7..8da6fdf3 100644 --- a/lib/webhooks.js +++ b/lib/webhooks.js @@ -1,11 +1,14 @@ /** * Get the default webhook signing secret * + * @param {object} [options] + * @param {AbortSignal} [options.signal] - An optional AbortSignal * @returns {Promise} Resolves with the signing secret for the default webhook */ -async function getDefaultWebhookSecret() { +async function getDefaultWebhookSecret({ signal } = {}) { const response = await this.request("/webhooks/default/secret", { method: "GET", + signal, }); return response.json(); diff --git a/package-lock.json b/package-lock.json index 300e5382..6c869967 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "replicate", - "version": "1.0.0", + "version": "1.4.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "replicate", - "version": "1.0.0", + "version": "1.4.0", "license": "Apache-2.0", "devDependencies": { "@biomejs/biome": "^1.4.1", diff --git a/package.json b/package.json index fcd47e8f..90cde175 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "replicate", - "version": "1.0.0", + "version": "1.4.0", "description": "JavaScript client for Replicate", "repository": "github:replicate/replicate-javascript", "homepage": "https://github.com/replicate/replicate-javascript#readme",