From ee62cc42aaaf98b3b210c1bf82b399a714412e88 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Mon, 8 Jan 2024 14:46:19 -1000 Subject: [PATCH 1/5] Add support for files API endpoints --- index.d.ts | 34 ++++++++++++++++++ index.js | 29 +++++++++++++++- index.test.ts | 40 +++++++++++++++++++-- lib/deployments.js | 4 +-- lib/files.js | 86 ++++++++++++++++++++++++++++++++++++++++++++++ lib/predictions.js | 6 ++-- lib/util.js | 28 ++++++++++++--- 7 files changed, 213 insertions(+), 14 deletions(-) create mode 100644 lib/files.js diff --git a/index.d.ts b/index.d.ts index 31a23251..6951247c 100644 --- a/index.d.ts +++ b/index.d.ts @@ -39,6 +39,21 @@ declare module "replicate" { }; } + export interface FileObject { + id: string; + name: string; + content_type: string; + size: number; + etag: string; + checksum: string; + metadata: Record; + created_at: string; + expires_at: string | null; + urls: { + get: string; + }; + } + export interface Hardware { sku: string; name: string; @@ -119,12 +134,16 @@ declare module "replicate" { input: Request | string, init?: RequestInit ) => Promise; + prepareInputs?: ( + input: Record + ) => Promise>; }); auth: string; userAgent?: string; baseUrl?: string; fetch: (input: Request | string, init?: RequestInit) => Promise; + prepareInputs: (input: Record) => Promise>; run( identifier: `${string}/${string}` | `${string}/${string}:${string}`, @@ -220,6 +239,16 @@ declare module "replicate" { list(): Promise>; }; + files: { + create: ( + file: File | Blob, + metadata?: Record + ) => Promise; + list: () => Promise; + get: (file_id: string) => Promise; + delete: (file_id: string) => Promise; + }; + hardware: { list(): Promise; }; @@ -310,4 +339,9 @@ declare module "replicate" { current: number; total: number; } | null; + + export function transformFileInputs( + inputs: Record, + options: { fallbackToDataURI: boolean } + ): Promise>; } diff --git a/index.js b/index.js index f4c0e2cf..566321cc 100644 --- a/index.js +++ b/index.js @@ -6,11 +6,14 @@ const { validateWebhook, parseProgressFromLogs, streamAsyncIterator, + transformFileInputsToReplicateFileURLs, + transformFileInputsToBase64EncodedDataURIs, } = require("./lib/util"); const accounts = require("./lib/accounts"); const collections = require("./lib/collections"); const deployments = require("./lib/deployments"); +const files = require("./lib/files"); const hardware = require("./lib/hardware"); const models = require("./lib/models"); const predictions = require("./lib/predictions"); @@ -46,6 +49,7 @@ class Replicate { * @param {string} options.userAgent - Identifier of your app * @param {string} [options.baseUrl] - Defaults to https://api.replicate.com/v1 * @param {Function} [options.fetch] - Fetch function to use. Defaults to `globalThis.fetch` + * @param {Function} [options.prepareInput] - Function to prepare input data before sending it to the API. */ constructor(options = {}) { this.auth = @@ -55,6 +59,15 @@ class Replicate { options.userAgent || `replicate-javascript/${packageJSON.version}`; this.baseUrl = options.baseUrl || "https://api.replicate.com/v1"; this.fetch = options.fetch || globalThis.fetch; + this.prepareInputs = + options.prepareInputs || + (async (inputs) => { + try { + return await transformFileInputsToReplicateFileURLs(this, inputs); + } catch (error) { + return await transformFileInputsToBase64EncodedDataURIs(inputs); + } + }); this.accounts = { current: accounts.current.bind(this), @@ -75,6 +88,13 @@ class Replicate { }, }; + this.files = { + create: files.create.bind(this), + list: files.list.bind(this), + get: files.get.bind(this), + delete: files.delete.bind(this), + }; + this.hardware = { list: hardware.list.bind(this), }; @@ -230,10 +250,17 @@ class Replicate { } } + let body = undefined; + if (data instanceof FormData) { + body = data; + } else if (data) { + body = JSON.stringify(data); + } + const init = { method, headers, - body: data ? JSON.stringify(data) : undefined, + body, }; const shouldRetry = diff --git a/index.test.ts b/index.test.ts index 53737e05..f879e25c 100644 --- a/index.test.ts +++ b/index.test.ts @@ -222,7 +222,7 @@ describe("Replicate client", () => { expect(prediction.id).toBe("ufawqhfynnddngldkgtslldrkq"); }); - test.each([ + const fileTestCases = [ // Skip test case if File type is not available ...(typeof File !== "undefined" ? [ @@ -245,11 +245,47 @@ describe("Replicate client", () => { value: Buffer.from("hello world"), expected: "data:application/octet-stream;base64,aGVsbG8gd29ybGQ=", }, - ])( + ]; + + test.each(fileTestCases)( + "converts a $type input into a Replicate file URL", + async ({ value: data, expected }) => { + nock(BASE_URL) + .post("/files") + .reply(201, { + urls: { + get: "https://replicate.com/api/files/123", + }, + }) + .post("/predictions") + .reply(201, (uri: string, body: Record) => { + return body; + }); + + const prediction = await client.predictions.create({ + version: + "5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa", + input: { + prompt: "Tell me a story", + data, + }, + stream: true, + }); + + expect(prediction.input).toEqual({ + prompt: "Tell me a story", + data: "https://replicate.com/api/files/123", + }); + } + ); + + test.each(fileTestCases)( "converts a $type input into a base64 encoded string", async ({ value: data, expected }) => { let actual: Record | undefined; nock(BASE_URL) + .post("/files") + .reply(503, "Service Unavailable") .post("/predictions") .reply(201, (_uri: string, body: Record) => { actual = body; diff --git a/lib/deployments.js b/lib/deployments.js index 4f6f3c6b..696c35fc 100644 --- a/lib/deployments.js +++ b/lib/deployments.js @@ -1,5 +1,3 @@ -const { transformFileInputs } = require("./util"); - /** * Create a new prediction with a deployment * @@ -30,7 +28,7 @@ async function createPrediction(deployment_owner, deployment_name, options) { method: "POST", data: { ...data, - input: await transformFileInputs(input), + input: await this.prepareInputs(input), stream, }, } diff --git a/lib/files.js b/lib/files.js new file mode 100644 index 00000000..f51bbd54 --- /dev/null +++ b/lib/files.js @@ -0,0 +1,86 @@ +/** + * Create a file + * + * @param {object} file - Required. The file object. + * @param {object} metadata - Optional. User-provided metadata associated with the file. + * @returns {Promise} - Resolves with the file data + */ +async function createFile(file, metadata = {}) { + const form = new FormData(); + + let filename; + let blob; + if (file instanceof Blob) { + fileName = file.name || `blob_${Date.now()}`; + blob = file; + } else if (Buffer.isBuffer(file)) { + fileName = `buffer_${Date.now()}`; + blob = new Blob(file, { type: "application/octet-stream" }); + } else { + throw new Error("Invalid file argument, must be a Blob or File"); + } + + form.append("content", blob, filename); + form.append( + "metadata", + new Blob([JSON.stringify(metadata)], { type: "application/json" }) + ); + + const response = await this.request("/files", { + method: "POST", + data: form, + headers: { + "Content-Type": "multipart/form-data", + }, + }); + + return response.json(); +} + +/** + * List all files + * + * @returns {Promise} - Resolves with the files data + */ +async function listFiles() { + const response = await this.request("/files", { + method: "GET", + }); + + return response.json(); +} + +/** + * Get a file + * + * @param {string} file_id - Required. The ID of the file. + * @returns {Promise} - Resolves with the file data + */ +async function getFile(file_id) { + const response = await this.request(`/files/${file_id}`, { + method: "GET", + }); + + return response.json(); +} + +/** + * Delete a file + * + * @param {string} file_id - Required. The ID of the file. + * @returns {Promise} - Resolves with the deletion confirmation + */ +async function deleteFile(file_id) { + const response = await this.request(`/files/${file_id}`, { + method: "DELETE", + }); + + return response.json(); +} + +module.exports = { + create: createFile, + list: listFiles, + get: getFile, + delete: deleteFile, +}; diff --git a/lib/predictions.js b/lib/predictions.js index 5b0370e5..ac9f5275 100644 --- a/lib/predictions.js +++ b/lib/predictions.js @@ -1,5 +1,3 @@ -const { transformFileInputs } = require("./util"); - /** * Create a new prediction * @@ -30,7 +28,7 @@ async function createPrediction(options) { method: "POST", data: { ...data, - input: await transformFileInputs(input), + input: await this.prepareInputs(input), version, stream, }, @@ -40,7 +38,7 @@ async function createPrediction(options) { method: "POST", data: { ...data, - input: await transformFileInputs(input), + input: await this.prepareInputs(input), stream, }, }); diff --git a/lib/util.js b/lib/util.js index 68b1d9de..b17843b6 100644 --- a/lib/util.js +++ b/lib/util.js @@ -215,6 +215,26 @@ async function withAutomaticRetries(request, options = {}) { return request(); } +/** + * Walks the inputs and, for any File or Blob, tries to upload it to Replicate + * and replaces the input with the URL of the uploaded file. + * + * @param {Replicate} client - The client used to upload the file + * @param {object} inputs - The inputs to transform + * @returns {object} - The transformed inputs + * @throws {ApiError} If the request to upload the file fails + */ +async function transformFileInputsToReplicateFileURLs(client, inputs) { + return await transform(inputs, async (value) => { + if (value instanceof Blob || value instanceof Buffer) { + const file = await client.files.create(value); + return file.urls.get; + } + + return value; + }); +} + const MAX_DATA_URI_SIZE = 10_000_000; /** @@ -225,9 +245,9 @@ const MAX_DATA_URI_SIZE = 10_000_000; * @returns {object} - The transformed inputs * @throws {Error} If the size of inputs exceeds a given threshould set by MAX_DATA_URI_SIZE */ -async function transformFileInputs(inputs) { +async function transformFileInputsToBase64EncodedDataURIs(inputs) { let totalBytes = 0; - const result = await transform(inputs, async (value) => { + return await transform(inputs, async (value) => { let buffer; let mime; @@ -258,8 +278,6 @@ async function transformFileInputs(inputs) { return `data:${mime};base64,${data}`; }); - - return result; } // Walk a JavaScript object and transform the leaf values. @@ -394,4 +412,6 @@ module.exports = { withAutomaticRetries, parseProgressFromLogs, streamAsyncIterator, + transformFileInputsToBase64EncodedDataURIs, + transformFileInputsToReplicateFileURLs, }; From 1bf586b670cf024271d3cc1e82355666d55aa180 Mon Sep 17 00:00:00 2001 From: Mattt Date: Thu, 22 Feb 2024 03:51:05 -0800 Subject: [PATCH 2/5] Apply suggestions from code review Co-authored-by: Aron Carroll --- index.d.ts | 2 +- lib/files.js | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/index.d.ts b/index.d.ts index 6951247c..96a52eef 100644 --- a/index.d.ts +++ b/index.d.ts @@ -46,7 +46,7 @@ declare module "replicate" { size: number; etag: string; checksum: string; - metadata: Record; + metadata: Record; created_at: string; expires_at: string | null; urls: { diff --git a/lib/files.js b/lib/files.js index f51bbd54..f6620e92 100644 --- a/lib/files.js +++ b/lib/files.js @@ -11,13 +11,13 @@ async function createFile(file, metadata = {}) { let filename; let blob; if (file instanceof Blob) { - fileName = file.name || `blob_${Date.now()}`; + filename = file.name || `blob_${Date.now()}`; blob = file; } else if (Buffer.isBuffer(file)) { - fileName = `buffer_${Date.now()}`; + filename = `buffer_${Date.now()}`; blob = new Blob(file, { type: "application/octet-stream" }); } else { - throw new Error("Invalid file argument, must be a Blob or File"); + throw new Error("Invalid file argument, must be a Blob, File or Buffer"); } form.append("content", blob, filename); From 37c1f74f0dc2a5359473e7d7d56c96f4ab87ad8d Mon Sep 17 00:00:00 2001 From: Aron Carroll Date: Tue, 5 Mar 2024 16:53:33 +0000 Subject: [PATCH 3/5] Remove prepareInputs in favor of FileEncodingStrategy This allows the user to determine if file uploads, falling back to base64 should be used or just sticking to one approach. The tests have been updated to validate the file upload payload and to ensure the url is correctly passed to the prediction create method. --- index.d.ts | 13 ++++--------- index.js | 12 ++---------- index.test.ts | 28 ++++++++++++++++++++++++---- lib/deployments.js | 8 +++++++- lib/predictions.js | 14 ++++++++++++-- lib/util.js | 28 +++++++++++++++++++++++++++- 6 files changed, 76 insertions(+), 27 deletions(-) diff --git a/index.d.ts b/index.d.ts index 96a52eef..f949afad 100644 --- a/index.d.ts +++ b/index.d.ts @@ -108,6 +108,8 @@ declare module "replicate" { export type Training = Prediction; + export type FileEncodingStrategy = "default" | "upload" | "data-uri"; + export interface Page { previous?: string; next?: string; @@ -134,16 +136,14 @@ declare module "replicate" { input: Request | string, init?: RequestInit ) => Promise; - prepareInputs?: ( - input: Record - ) => Promise>; + fileEncodingStrategy?: FileEncodingStrategy; }); auth: string; userAgent?: string; baseUrl?: string; fetch: (input: Request | string, init?: RequestInit) => Promise; - prepareInputs: (input: Record) => Promise>; + fileEncodingStrategy: FileEncodingStrategy; run( identifier: `${string}/${string}` | `${string}/${string}:${string}`, @@ -339,9 +339,4 @@ declare module "replicate" { current: number; total: number; } | null; - - export function transformFileInputs( - inputs: Record, - options: { fallbackToDataURI: boolean } - ): Promise>; } diff --git a/index.js b/index.js index 566321cc..d2568a2e 100644 --- a/index.js +++ b/index.js @@ -49,7 +49,7 @@ class Replicate { * @param {string} options.userAgent - Identifier of your app * @param {string} [options.baseUrl] - Defaults to https://api.replicate.com/v1 * @param {Function} [options.fetch] - Fetch function to use. Defaults to `globalThis.fetch` - * @param {Function} [options.prepareInput] - Function to prepare input data before sending it to the API. + * @param {"default" | "upload" | "data-uri"} [options.fileEncodingStrategy] - Determines the file encoding strategy to use */ constructor(options = {}) { this.auth = @@ -59,15 +59,7 @@ class Replicate { options.userAgent || `replicate-javascript/${packageJSON.version}`; this.baseUrl = options.baseUrl || "https://api.replicate.com/v1"; this.fetch = options.fetch || globalThis.fetch; - this.prepareInputs = - options.prepareInputs || - (async (inputs) => { - try { - return await transformFileInputsToReplicateFileURLs(this, inputs); - } catch (error) { - return await transformFileInputsToBase64EncodedDataURIs(inputs); - } - }); + this.fileEncodingStrategy = options.fileEncodingStrategy ?? "default"; this.accounts = { current: accounts.current.bind(this), diff --git a/index.test.ts b/index.test.ts index f879e25c..7502969d 100644 --- a/index.test.ts +++ b/index.test.ts @@ -228,7 +228,7 @@ describe("Replicate client", () => { ? [ { type: "file", - value: new File(["hello world"], "hello.txt", { + value: new File(["hello world"], "file_hello.txt", { type: "text/plain", }), expected: "data:text/plain;base64,aGVsbG8gd29ybGQ=", @@ -249,16 +249,22 @@ describe("Replicate client", () => { test.each(fileTestCases)( "converts a $type input into a Replicate file URL", - async ({ value: data, expected }) => { + async ({ value: data, type }) => { + const mockedFetch = jest.spyOn(client, "fetch"); + nock(BASE_URL) .post("/files") + .matchHeader("Content-Type", "multipart/form-data") .reply(201, { urls: { get: "https://replicate.com/api/files/123", }, }) - .post("/predictions") - .reply(201, (uri: string, body: Record) => { + .post( + "/predictions", + (body) => body.input.data === "https://replicate.com/api/files/123" + ) + .reply(201, (_uri: string, body: Record) => { return body; }); @@ -272,6 +278,20 @@ describe("Replicate client", () => { stream: true, }); + expect(client.fetch).toHaveBeenCalledWith( + new URL("https://api.replicate.com/v1/files"), + { + method: "POST", + body: expect.any(FormData), + headers: expect.objectContaining({ + "Content-Type": "multipart/form-data", + }), + } + ); + const form = mockedFetch.mock.calls[0][1]?.body as FormData; + // @ts-ignore + expect(form?.get("content")?.name).toMatch(new RegExp(`^${type}_`)); + expect(prediction.input).toEqual({ prompt: "Tell me a story", data: "https://replicate.com/api/files/123", diff --git a/lib/deployments.js b/lib/deployments.js index 696c35fc..27a2f6a2 100644 --- a/lib/deployments.js +++ b/lib/deployments.js @@ -1,3 +1,5 @@ +const { transformFileInputs } = require("./util"); + /** * Create a new prediction with a deployment * @@ -28,7 +30,11 @@ async function createPrediction(deployment_owner, deployment_name, options) { method: "POST", data: { ...data, - input: await this.prepareInputs(input), + input: await transformFileInputs( + this, + input, + this.fileEncodingStrategy + ), stream, }, } diff --git a/lib/predictions.js b/lib/predictions.js index ac9f5275..c290d40f 100644 --- a/lib/predictions.js +++ b/lib/predictions.js @@ -1,3 +1,5 @@ +const { transformFileInputs } = require("./util"); + /** * Create a new prediction * @@ -28,7 +30,11 @@ async function createPrediction(options) { method: "POST", data: { ...data, - input: await this.prepareInputs(input), + input: await transformFileInputs( + this, + input, + this.fileEncodingStrategy + ), version, stream, }, @@ -38,7 +44,11 @@ async function createPrediction(options) { method: "POST", data: { ...data, - input: await this.prepareInputs(input), + input: await transformFileInputs( + this, + input, + this.fileEncodingStrategy + ), stream, }, }); diff --git a/lib/util.js b/lib/util.js index b17843b6..dfec202b 100644 --- a/lib/util.js +++ b/lib/util.js @@ -209,12 +209,38 @@ async function withAutomaticRetries(request, options = {}) { } attempts += 1; } - /* eslint-enable no-await-in-loop */ } while (attempts < maxRetries); return request(); } +/** + * Walks the inputs and, for any File or Blob, tries to upload it to Replicate + * and replaces the input with the URL of the uploaded file. + * + * @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 + * @throws {ApiError} If the request to upload the file fails + */ +async function transformFileInputs(client, inputs, strategy) { + switch (strategy) { + case "data-uri": + return await transformFileInputsToBase64EncodedDataURIs(client, inputs); + case "upload": + return await transformFileInputsToReplicateFileURLs(client, inputs); + case "default": + try { + return await transformFileInputsToReplicateFileURLs(client, inputs); + } catch (error) { + return await transformFileInputsToBase64EncodedDataURIs(inputs); + } + default: + throw new Error(`Unexpected file upload strategy: ${strategy}`); + } +} + /** * Walks the inputs and, for any File or Blob, tries to upload it to Replicate * and replaces the input with the URL of the uploaded file. From fcdcaa33e8827f83caf1113bef7652d2b2270193 Mon Sep 17 00:00:00 2001 From: Aron Carroll Date: Mon, 11 Mar 2024 11:54:18 +0000 Subject: [PATCH 4/5] Remove replicate.files API --- index.d.ts | 10 ---------- index.js | 10 ---------- lib/util.js | 5 ++--- 3 files changed, 2 insertions(+), 23 deletions(-) diff --git a/index.d.ts b/index.d.ts index f949afad..1ef9e89c 100644 --- a/index.d.ts +++ b/index.d.ts @@ -239,16 +239,6 @@ declare module "replicate" { list(): Promise>; }; - files: { - create: ( - file: File | Blob, - metadata?: Record - ) => Promise; - list: () => Promise; - get: (file_id: string) => Promise; - delete: (file_id: string) => Promise; - }; - hardware: { list(): Promise; }; diff --git a/index.js b/index.js index d2568a2e..3ef0d3d8 100644 --- a/index.js +++ b/index.js @@ -6,14 +6,11 @@ const { validateWebhook, parseProgressFromLogs, streamAsyncIterator, - transformFileInputsToReplicateFileURLs, - transformFileInputsToBase64EncodedDataURIs, } = require("./lib/util"); const accounts = require("./lib/accounts"); const collections = require("./lib/collections"); const deployments = require("./lib/deployments"); -const files = require("./lib/files"); const hardware = require("./lib/hardware"); const models = require("./lib/models"); const predictions = require("./lib/predictions"); @@ -80,13 +77,6 @@ class Replicate { }, }; - this.files = { - create: files.create.bind(this), - list: files.list.bind(this), - get: files.get.bind(this), - delete: files.delete.bind(this), - }; - this.hardware = { list: hardware.list.bind(this), }; diff --git a/lib/util.js b/lib/util.js index dfec202b..e164899e 100644 --- a/lib/util.js +++ b/lib/util.js @@ -1,4 +1,5 @@ const ApiError = require("./error"); +const { create: createFile } = require("./files"); /** * @see {@link validateWebhook} @@ -253,7 +254,7 @@ async function transformFileInputs(client, inputs, strategy) { async function transformFileInputsToReplicateFileURLs(client, inputs) { return await transform(inputs, async (value) => { if (value instanceof Blob || value instanceof Buffer) { - const file = await client.files.create(value); + const file = await createFile.call(client, value); return file.urls.get; } @@ -438,6 +439,4 @@ module.exports = { withAutomaticRetries, parseProgressFromLogs, streamAsyncIterator, - transformFileInputsToBase64EncodedDataURIs, - transformFileInputsToReplicateFileURLs, }; From e429fcdda46307935a1ed5d62cf47e7880c5e018 Mon Sep 17 00:00:00 2001 From: Aron Carroll Date: Mon, 11 Mar 2024 11:56:43 +0000 Subject: [PATCH 5/5] Tell Biome to ignore .wrangler directory --- biome.json | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/biome.json b/biome.json index ecb665f5..094cf0ec 100644 --- a/biome.json +++ b/biome.json @@ -1,7 +1,11 @@ { "$schema": "https://biomejs.dev/schemas/1.0.0/schema.json", "files": { - "ignore": [".wrangler", "vendor/*"] + "ignore": [ + ".wrangler", + "node_modules", + "vendor/*" + ] }, "formatter": { "indentStyle": "space",