Merge branch 'main' into mattt/deprecate-stream-parameter

Author: mattt

.github/workflows/ci.yml

@@ -14,7 +14,7 @@ jobs:
       matrix:
         suite: [node]
         # See supported Node.js release schedule at https://nodejs.org/en/about/previous-releases
-        node-version: [18.x, 20.x, 22.4] # TODO: unpin to 22.x once https://github.com/nodejs/node/issues/53902 is resolved
+        node-version: [18.x, 20.x, 22.x]
 
     steps:
       - uses: actions/checkout@v4

README.md

@@ -455,6 +455,18 @@ const response = await replicate.models.list();
 }
 ```
 
+### `replicate.models.search`
+
+Search for public models on Replicate.
+
+```js
+const response = await replicate.models.search(query);
+```
+
+| name    | type   | description                            |
+| ------- | ------ | -------------------------------------- |
+| `query` | string | **Required**. The search query string. |
+
 ### `replicate.models.create`
 
 Create a new public or private model.

index.d.ts

@@ -8,6 +8,12 @@ declare module "replicate" {
     response: Response;
   }
 
+  export interface FileOutput extends ReadableStream {
+    blob(): Promise<Blob>;
+    url(): URL;
+    toString(): string;
+  }
+
   export interface Account {
     type: "user" | "organization";
     username: string;
@@ -137,6 +143,7 @@ declare module "replicate" {
         init?: RequestInit
       ) => Promise<Response>;
       fileEncodingStrategy?: FileEncodingStrategy;
+      useFileOutput?: boolean;
     });
 
     auth: string;
@@ -281,6 +288,7 @@ declare module "replicate" {
           version_id: string
         ): Promise<ModelVersion>;
       };
+      search(query: string): Promise<Page<Model>>;
     };
 
     predictions: {

index.js

@@ -1,7 +1,8 @@
 const ApiError = require("./lib/error");
 const ModelVersionIdentifier = require("./lib/identifier");
-const { createReadableStream } = require("./lib/stream");
+const { createReadableStream, createFileOutput } = require("./lib/stream");
 const {
+  transform,
   withAutomaticRetries,
   validateWebhook,
   parseProgressFromLogs,
@@ -47,6 +48,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 {boolean} [options.useFileOutput] - Set to `true` to return `FileOutput` objects from `run` instead of URLs, defaults to false.
    * @param {"default" | "upload" | "data-uri"} [options.fileEncodingStrategy] - Determines the file encoding strategy to use
    */
   constructor(options = {}) {
@@ -57,7 +59,8 @@ class Replicate {
       options.userAgent || `replicate-javascript/${packageJSON.version}`;
     this.baseUrl = options.baseUrl || "https://api.replicate.com/v1";
     this.fetch = options.fetch || globalThis.fetch;
-    this.fileEncodingStrategy = options.fileEncodingStrategy ?? "default";
+    this.fileEncodingStrategy = options.fileEncodingStrategy || "default";
+    this.useFileOutput = options.useFileOutput || false;
 
     this.accounts = {
       current: accounts.current.bind(this),
@@ -98,6 +101,7 @@ class Replicate {
         list: models.versions.list.bind(this),
         get: models.versions.get.bind(this),
       },
+      search: models.search.bind(this),
     };
 
     this.predictions = {
@@ -195,7 +199,17 @@ class Replicate {
       throw new Error(`Prediction failed: ${prediction.error}`);
     }
 
-    return prediction.output;
+    return transform(prediction.output, (value) => {
+      if (
+        typeof value === "string" &&
+        (value.startsWith("https:") || value.startsWith("data:"))
+      ) {
+        return this.useFileOutput
+          ? createFileOutput({ url: value, fetch: this.fetch })
+          : value;
+      }
+      return value;
+    });
   }
 
   /**

index.test.ts

@@ -1,13 +1,13 @@
 import { expect, jest, test } from "@jest/globals";
 import Replicate, {
   ApiError,
+  FileOutput,
   Model,
   Prediction,
   validateWebhook,
   parseProgressFromLogs,
 } from "replicate";
 import nock from "nock";
-import { Readable } from "node:stream";
 import { createReadableStream } from "./lib/stream";
 
 let client: Replicate;
@@ -1053,7 +1053,7 @@ describe("Replicate client", () => {
   describe("predictions.create with model", () => {
     test("Calls the correct API route with the correct payload", async () => {
       nock(BASE_URL)
-        .post("/models/meta/llama-2-70b-chat/predictions")
+        .post("/models/meta/meta-llama-3-70b-instruct/predictions")
         .reply(200, {
           id: "heat2o3bzn3ahtr6bjfftvbaci",
           model: "replicate/lifeboat-70b",
@@ -1072,7 +1072,7 @@ describe("Replicate client", () => {
           },
         });
       const prediction = await client.predictions.create({
-        model: "meta/llama-2-70b-chat",
+        model: "meta/meta-llama-3-70b-instruct",
         input: {
           prompt: "Please write a haiku about llamas.",
         },
@@ -1217,6 +1217,44 @@ describe("Replicate client", () => {
     });
   });
 
+  describe("models.search", () => {
+    test("Calls the correct API route with the correct payload", async () => {
+      nock(BASE_URL)
+        .intercept("/models", "QUERY")
+        .reply(200, {
+          results: [
+            {
+              url: "https://replicate.com/meta/meta-llama-3-70b-instruct",
+              owner: "meta",
+              name: "meta-llama-3-70b-instruct",
+              description:
+                "Llama 2 is a collection of pretrained and fine-tuned generative text models ranging in scale from 7 billion to 70 billion parameters.",
+              visibility: "public",
+              github_url: null,
+              paper_url:
+                "https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/",
+              license_url: "https://ai.meta.com/llama/license/",
+              run_count: 1000000,
+              cover_image_url:
+                "https://replicate.delivery/pbxt/IJqFrnAKEDiCBnlXyndzVVxkZvfQ7kLjGVEZZPXTRXxOOPkQA/llama2.png",
+              default_example: null,
+              latest_version: null,
+            },
+            // ... more results ...
+          ],
+          next: null,
+          previous: null,
+        });
+
+      const searchResults = await client.models.search("llama");
+      expect(searchResults.results.length).toBeGreaterThan(0);
+      expect(searchResults.results[0].owner).toBe("meta");
+      expect(searchResults.results[0].name).toBe("meta-llama-3-70b-instruct");
+    });
+
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   describe("run", () => {
     test("Calls the correct API routes", async () => {
       nock(BASE_URL)
@@ -1504,6 +1542,203 @@ describe("Replicate client", () => {
 
       scope.done();
     });
+
+    test("returns FileOutput for URLs when useFileOutput is true", async () => {
+      client = new Replicate({ auth: "foo", useFileOutput: true });
+
+      nock(BASE_URL)
+        .post("/predictions")
+        .reply(201, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "starting",
+          logs: null,
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "succeeded",
+          output: "https://example.com",
+          logs: [].join("\n"),
+        });
+
+      nock("https://example.com")
+        .get("/")
+        .reply(200, "hello world", { "Content-Type": "text/plain" });
+
+      const output = (await client.run(
+        "owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
+        {
+          input: { text: "Hello, world!" },
+        }
+      )) as FileOutput;
+
+      expect(output).toBeInstanceOf(ReadableStream);
+      expect(output.url()).toEqual(new URL("https://example.com"));
+
+      const blob = await output.blob();
+      expect(blob.type).toEqual("text/plain");
+      expect(blob.arrayBuffer()).toEqual(
+        new Blob(["Hello, world!"]).arrayBuffer()
+      );
+    });
+
+    test("returns FileOutput for URLs when useFileOutput is true - acts like string", async () => {
+      client = new Replicate({ auth: "foo", useFileOutput: true });
+
+      nock(BASE_URL)
+        .post("/predictions")
+        .reply(201, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "starting",
+          logs: null,
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "succeeded",
+          output: "https://example.com",
+          logs: [].join("\n"),
+        });
+
+      nock("https://example.com")
+        .get("/")
+        .reply(200, "hello world", { "Content-Type": "text/plain" });
+
+      const output = (await client.run(
+        "owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
+        {
+          input: { text: "Hello, world!" },
+        }
+      )) as unknown as string;
+
+      expect(fetch(output).then((r) => r.text())).resolves.toEqual(
+        "hello world"
+      );
+    });
+
+    test("returns FileOutput for URLs when useFileOutput is true - array output", async () => {
+      client = new Replicate({ auth: "foo", useFileOutput: true });
+
+      nock(BASE_URL)
+        .post("/predictions")
+        .reply(201, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "starting",
+          logs: null,
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "succeeded",
+          output: ["https://example.com"],
+          logs: [].join("\n"),
+        });
+
+      nock("https://example.com")
+        .get("/")
+        .reply(200, "hello world", { "Content-Type": "text/plain" });
+
+      const [output] = (await client.run(
+        "owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
+        {
+          input: { text: "Hello, world!" },
+        }
+      )) as FileOutput[];
+
+      expect(output).toBeInstanceOf(ReadableStream);
+      expect(output.url()).toEqual(new URL("https://example.com"));
+
+      const blob = await output.blob();
+      expect(blob.type).toEqual("text/plain");
+      expect(blob.arrayBuffer()).toEqual(
+        new Blob(["Hello, world!"]).arrayBuffer()
+      );
+    });
+
+    test("returns FileOutput for URLs when useFileOutput is true - data uri", async () => {
+      client = new Replicate({ auth: "foo", useFileOutput: true });
+
+      nock(BASE_URL)
+        .post("/predictions")
+        .reply(201, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "starting",
+          logs: null,
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "processing",
+          logs: [].join("\n"),
+        })
+        .get("/predictions/ufawqhfynnddngldkgtslldrkq")
+        .reply(200, {
+          id: "ufawqhfynnddngldkgtslldrkq",
+          status: "succeeded",
+          output: "data:text/plain;base64,SGVsbG8sIHdvcmxkIQ==",
+          logs: [].join("\n"),
+        });
+
+      const output = (await client.run(
+        "owner/model:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
+        {
+          input: { text: "Hello, world!" },
+        }
+      )) as FileOutput;
+
+      expect(output).toBeInstanceOf(ReadableStream);
+      expect(output.url()).toEqual(
+        new URL("data:text/plain;base64,SGVsbG8sIHdvcmxkIQ==")
+      );
+
+      const blob = await output.blob();
+      expect(blob.type).toEqual("text/plain");
+      expect(blob.arrayBuffer()).toEqual(
+        new Blob(["Hello, world!"]).arrayBuffer()
+      );
+    });
   });
 
   describe("webhooks.default.secret.get", () => {