diff --git a/index.d.ts b/index.d.ts index c415bf0..1218cf6 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 0552ea5..8736bb8 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,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} - 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 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); diff --git a/index.test.ts b/index.test.ts index ec9e523..2684b85 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') @@ -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(); diff --git a/lib/identifier.js b/lib/identifier.js new file mode 100644 index 0000000..07e21d1 --- /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;