From b67c709e87dfee19c79d3cba633f6c05d7580baf Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 8 Dec 2023 04:57:33 -0800 Subject: [PATCH 1/2] Extract model version identifier into separate component --- index.js | 28 ++++++---------------------- index.test.ts | 2 -- lib/identifier.js | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 24 deletions(-) create mode 100644 lib/identifier.js diff --git a/index.js b/index.js index 0552ea51..ebe91dc9 100644 --- a/index.js +++ b/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,35 +101,18 @@ 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} - 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( - `^(?${namePattern.source})/(?${namePattern.source}):(?[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 { version } = match.groups; + const identifier = ModelVersionIdentifier.parse(ref); let prediction = await this.predictions.create({ ...data, - version, + version: identifier.version, }); // Call progress callback with the initial prediction object diff --git a/index.test.ts b/index.test.ts index ec9e523d..7c33d6dd 100644 --- a/index.test.ts +++ b/index.test.ts @@ -828,8 +828,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(); diff --git a/lib/identifier.js b/lib/identifier.js new file mode 100644 index 00000000..07e21d1e --- /dev/null +++ b/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(/^(?[^/]+)\/(?[^/:]+)(:(?.+))?$/); + 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; From 4aa4dabf54a9b96bcd15b013db7e75aa0b2a2a53 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Fri, 8 Dec 2023 05:07:00 -0800 Subject: [PATCH 2/2] Allow `run` method to take model argument, when supported --- index.d.ts | 2 +- index.js | 17 ++++++++++---- index.test.ts | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/index.d.ts b/index.d.ts index c415bf08..1218cf68 100644 --- a/index.d.ts +++ b/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 }; diff --git a/index.js b/index.js index ebe91dc9..8736bb8f 100644 --- a/index.js +++ b/index.js @@ -110,10 +110,19 @@ class Replicate { const identifier = ModelVersionIdentifier.parse(ref); - let prediction = await this.predictions.create({ - ...data, - version: identifier.version, - }); + 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 + ); + } // Call progress callback with the initial prediction object if (progress) { diff --git a/index.test.ts b/index.test.ts index 7c33d6dd..2684b85f 100644 --- a/index.test.ts +++ b/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')