Add support for files API endpoints

Author: mattt

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<string, any>;
+    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<Response>;
+      prepareInputs?: (
+        input: Record<string, any>
+      ) => Promise<Record<string, any>>;
     });
 
     auth: string;
     userAgent?: string;
     baseUrl?: string;
     fetch: (input: Request | string, init?: RequestInit) => Promise<Response>;
+    prepareInputs: (input: Record<string, any>) => Promise<Record<string, any>>;
 
     run(
       identifier: `${string}/${string}` | `${string}/${string}:${string}`,
@@ -196,6 +215,16 @@ declare module "replicate" {
       ): Promise<Deployment>;
     };
 
+    files: {
+      create: (
+        file: File | Blob,
+        metadata?: Record<string, string | number | boolean | null>
+      ) => Promise<FileObject>;
+      list: () => Promise<FileObject>;
+      get: (file_id: string) => Promise<FileObject>;
+      delete: (file_id: string) => Promise<FileObject>;
+    };
+
     hardware: {
       list(): Promise<Hardware[]>;
     };
@@ -286,4 +315,9 @@ declare module "replicate" {
     current: number;
     total: number;
   } | null;
+
+  export function transformFileInputs(
+    inputs: Record<string, any>,
+    options: { fallbackToDataURI: boolean }
+  ): Promise<Record<string, any>>;
 }

index.js

@@ -5,11 +5,14 @@ const {
   withAutomaticRetries,
   validateWebhook,
   parseProgressFromLogs,
+  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");
@@ -45,13 +48,23 @@ 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 = options.auth || process.env.REPLICATE_API_TOKEN;
     this.userAgent =
       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),
@@ -69,6 +82,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),
     };
@@ -222,10 +242,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 =

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<string, any>) => {
+            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<string, any> | undefined;
         nock(BASE_URL)
+          .post("/files")
+          .reply(503, "Service Unavailable")
           .post("/predictions")
           .reply(201, (uri: string, body: Record<string, any>) => {
             actual = body;

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,
       },
     }

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<object>} - 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<object>} - 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<object>} - 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<object>} - 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,
+};

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,
       },
     });

lib/util.js

@@ -208,6 +208,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;
 
 /**
@@ -218,9 +238,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;
 
@@ -251,8 +271,6 @@ async function transformFileInputs(inputs) {
 
     return `data:${mime};base64,${data}`;
   });
-
-  return result;
 }
 
 // Walk a JavaScript object and transform the leaf values.
@@ -349,4 +367,6 @@ module.exports = {
   validateWebhook,
   withAutomaticRetries,
   parseProgressFromLogs,
+  transformFileInputsToBase64EncodedDataURIs,
+  transformFileInputsToReplicateFileURLs,
 };