Allow run method to take model argument, when supported (#167)

* Extract model version identifier into separate component

* Allow `run` method to take model argument, when supported

Author: mattt

index.d.ts

@@ -89,7 +89,7 @@ declare module 'replicate' {
     fetch: Function;
 
     run(
-      identifier: `${string}/${string}:${string}`,
+      identifier: `${string}/${string}` | `${string}/${string}:${string}`,
       options: {
         input: object;
         wait?: { interval?: number };

index.js

@@ -1,4 +1,5 @@
 const ApiError = require('./lib/error');
+const ModelVersionIdentifier = require('./lib/identifier');
 const { withAutomaticRetries } = require('./lib/util');
 
 const collections = require('./lib/collections');
@@ -91,7 +92,7 @@ class Replicate {
   /**
    * Run a model and wait for its output.
    *
-   * @param {string} identifier - Required. The model version identifier in the format "{owner}/{name}:{version}"
+   * @param {string} ref - Required. The model version identifier in the format "owner/name" or "owner/name:version"
    * @param {object} options
    * @param {object} options.input - Required. An object with the model inputs
    * @param {object} [options.wait] - Options for waiting for the prediction to finish
@@ -100,37 +101,29 @@ class Replicate {
    * @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] - AbortSignal to cancel the prediction
    * @param {Function} [progress] - Callback function that receives the prediction object as it's updated. The function is called when the prediction is created, each time its updated while polling for completion, and when it's completed.
+   * @throws {Error} If the reference is invalid
    * @throws {Error} If the prediction failed
    * @returns {Promise<object>} - Resolves with the output of running the model
    */
-  async run(identifier, options, progress) {
+  async run(ref, options, progress) {
     const { wait, ...data } = options;
 
-    // Define a pattern for owner and model names that allows
-    // letters, digits, and certain special characters.
-    // Example: "user123", "abc__123", "user.name"
-    const namePattern = /[a-zA-Z0-9]+(?:(?:[._]|__|[-]*)[a-zA-Z0-9]+)*/;
-
-    // Define a pattern for "owner/name:version" format with named capturing groups.
-    // Example: "user123/repo_a:1a2b3c"
-    const pattern = new RegExp(
-      `^(?<owner>${namePattern.source})/(?<name>${namePattern.source}):(?<version>[0-9a-fA-F]+)$`
-    );
-
-    const match = identifier.match(pattern);
-    if (!match || !match.groups) {
-      throw new Error(
-        'Invalid version. It must be in the format "owner/name:version"'
+    const identifier = ModelVersionIdentifier.parse(ref);
+
+    let prediction;
+    if (identifier.version) {
+      prediction = await this.predictions.create({
+        ...data,
+        version: identifier.version,
+      });
+    } else {
+      prediction = await this.models.predictions.create(
+        identifier.owner,
+        identifier.name,
+        data
       );
     }
 
-    const { version } = match.groups;
-
-    let prediction = await this.predictions.create({
-      ...data,
-      version,
-    });
-
     // Call progress callback with the initial prediction object
     if (progress) {
       progress(prediction);

index.test.ts

@@ -749,7 +749,7 @@ describe('Replicate client', () => {
   });
 
   describe('run', () => {
-    test('Calls the correct API routes', async () => {
+    test('Calls the correct API routes for a version', async () => {
       let firstPollingRequest = true;
 
       nock(BASE_URL)
@@ -808,6 +808,65 @@ describe('Replicate client', () => {
       expect(progress).toHaveBeenCalledTimes(4);
     });
 
+    test('Calls the correct API routes for a model', async () => {
+      let firstPollingRequest = true;
+
+      nock(BASE_URL)
+        .post('/models/replicate/hello-world/predictions')
+        .reply(201, {
+          id: 'ufawqhfynnddngldkgtslldrkq',
+          status: 'starting',
+        })
+        .get('/predictions/ufawqhfynnddngldkgtslldrkq')
+        .twice()
+        .reply(200, {
+          id: 'ufawqhfynnddngldkgtslldrkq',
+          status: 'processing',
+        })
+        .get('/predictions/ufawqhfynnddngldkgtslldrkq')
+        .reply(200, {
+          id: 'ufawqhfynnddngldkgtslldrkq',
+          status: 'succeeded',
+          output: 'Goodbye!',
+        });
+
+      const progress = jest.fn();
+
+      const output = await client.run(
+        'replicate/hello-world',
+        {
+          input: { text: 'Hello, world!' },
+          wait: { interval: 1 }
+        },
+        progress
+      );
+
+      expect(output).toBe('Goodbye!');
+
+      expect(progress).toHaveBeenNthCalledWith(1, {
+        id: 'ufawqhfynnddngldkgtslldrkq',
+        status: 'starting',
+      });
+
+      expect(progress).toHaveBeenNthCalledWith(2, {
+        id: 'ufawqhfynnddngldkgtslldrkq',
+        status: 'processing',
+      });
+
+      expect(progress).toHaveBeenNthCalledWith(3, {
+        id: 'ufawqhfynnddngldkgtslldrkq',
+        status: 'processing',
+      });
+
+      expect(progress).toHaveBeenNthCalledWith(4, {
+        id: 'ufawqhfynnddngldkgtslldrkq',
+        status: 'succeeded',
+        output: 'Goodbye!',
+      });
+
+      expect(progress).toHaveBeenCalledTimes(4);
+    });
+
     test('Does not throw an error for identifier containing hyphen and full stop', async () => {
       nock(BASE_URL)
         .post('/predictions')
@@ -828,8 +887,6 @@ describe('Replicate client', () => {
     test('Throws an error for invalid identifiers', async () => {
       const options = { input: { text: 'Hello, world!' } }
 
-      await expect(client.run('owner/model:invalid', options)).rejects.toThrow();
-
       // @ts-expect-error
       await expect(client.run('owner:abc123', options)).rejects.toThrow();
 

lib/identifier.js

@@ -0,0 +1,35 @@
+/*
+ * A reference to a model version in the format `owner/name` or `owner/name:version`.
+ */
+class ModelVersionIdentifier {
+  /*
+   * @param {string} Required. The model owner.
+   * @param {string} Required. The model name.
+   * @param {string} The model version.
+   */
+  constructor(owner, name, version = null) {
+    this.owner = owner;
+    this.name = name;
+    this.version = version;
+  }
+
+  /*
+   * Parse a reference to a model version
+   *
+   * @param {string}
+   * @returns {ModelVersionIdentifier}
+   * @throws {Error} If the reference is invalid.
+   */
+  static parse(ref) {
+    const match = ref.match(/^(?<owner>[^/]+)\/(?<name>[^/:]+)(:(?<version>.+))?$/);
+    if (!match) {
+      throw new Error(`Invalid reference to model version: ${ref}. Expected format: owner/name or owner/name:version`);
+    }
+
+    const { owner, name, version } = match.groups;
+
+    return new ModelVersionIdentifier(owner, name, version);
+  }
+}
+
+module.exports = ModelVersionIdentifier;