From e4c2faa1eef7f3caa44a95c164bfb2d9c332d689 Mon Sep 17 00:00:00 2001 From: Santiago Bosio Date: Fri, 27 Sep 2024 19:02:16 -0300 Subject: [PATCH] Implemented 'ai:models:call' for image generation --- README.md | 11 +- src/commands/ai/docs.ts | 35 +--- src/commands/ai/models/call.ts | 62 +++++-- src/lib/ai/types.ts | 22 +++ src/lib/open-url.ts | 36 +++++ test/commands/ai/docs.test.ts | 131 +++------------ test/commands/ai/models/call.test.ts | 232 ++++++++++++++++++++++++++- test/helpers/fixtures.ts | 57 ++++++- test/lib/open-url.test.ts | 159 ++++++++++++++++++ 9 files changed, 583 insertions(+), 162 deletions(-) create mode 100644 src/lib/open-url.ts create mode 100644 test/lib/open-url.test.ts diff --git a/README.md b/README.md index d867129..1644333 100644 --- a/README.md +++ b/README.md @@ -101,18 +101,19 @@ make an inference request to a specific AI model resource ``` USAGE - $ heroku ai:models:call [MODEL_RESOURCE] -a -p [-j] [--optfile ] [--opts ] [-o - ] [-r ] + $ heroku ai:models:call [MODEL_RESOURCE] -p [-a ] [--browser ] [-j] [--optfile ] + [--opts ] [-o ] [-r ] ARGUMENTS MODEL_RESOURCE The resource ID or alias of the model to call. FLAGS - -a, --app= (required) app to run command against + -a, --app= app to run command against -j, --json Output response as JSON -o, --output= The file path where the command writes the model response. -p, --prompt= (required) The input prompt for the model. -r, --remote= git remote of app to use + --browser= browser to open images with (example: "firefox", "safari") --optfile= Additional options for model inference, provided as a JSON config file. --opts= Additional options for model inference, provided as a JSON string. @@ -120,9 +121,9 @@ DESCRIPTION make an inference request to a specific AI model resource EXAMPLES - $ heroku ai:models:call my_llm --prompt "What is the meaning of life?" + $ heroku ai:models:call my_llm --app my-app --prompt "What is the meaning of life?" - $ heroku ai:models:call sdxl --prompt "Generate an image of a sunset" --opts '{"quality": "hd"}' + $ heroku ai:models:call sdxl --app my-app --prompt "Generate an image of a sunset" --opts '{"quality":"hd"}' -o sunset.png ``` _See code: [dist/commands/ai/models/call.ts](https://github.com/heroku/heroku-cli-plugin-integration/blob/v0.0.0/dist/commands/ai/models/call.ts)_ diff --git a/src/commands/ai/docs.ts b/src/commands/ai/docs.ts index 24790eb..a4f04cc 100644 --- a/src/commands/ai/docs.ts +++ b/src/commands/ai/docs.ts @@ -1,8 +1,5 @@ -import color from '@heroku-cli/color' import {flags} from '@heroku-cli/command' -import {ux} from '@oclif/core' -import {CLIError} from '@oclif/core/lib/errors' -import open from 'open' +import {openUrl} from '../../lib/open-url' import Command from '../../lib/base' export default class Docs extends Command { @@ -12,39 +9,11 @@ export default class Docs extends Command { browser: flags.string({description: 'browser to open docs with (example: "firefox", "safari")'}), } - static urlOpener: (...args: Parameters) => ReturnType = open - public async run(): Promise { const {flags} = await this.parse(Docs) const browser = flags.browser const url = process.env.HEROKU_AI_DOCS_URL || Docs.defaultUrl - let browserErrorShown = false - const showBrowserError = (browser?: string) => { - if (browserErrorShown) return - - ux.warn(`Unable to open ${browser ? browser : 'your default'} browser. Please visit ${color.cyan(url)} to view the documentation.`) - browserErrorShown = true - } - - ux.log(`Opening ${color.cyan(url)} in ${browser ? browser : 'your default'} browser…`) - - try { - await ux.anykey( - `Press any key to open up the browser to show Heroku AI documentation, or ${color.yellow('q')} to exit` - ) - } catch (error) { - const {message, oclif} = error as CLIError - ux.error(message, {exit: oclif?.exit || 1}) - } - - const cp = await Docs.urlOpener(url, {wait: false, ...(browser ? {app: {name: browser}} : {})}) - cp.on('error', (err: Error) => { - ux.warn(err) - showBrowserError(browser) - }) - cp.on('close', (code: number) => { - if (code !== 0) showBrowserError(browser) - }) + await openUrl(url, browser, 'view the documentation') } } diff --git a/src/commands/ai/models/call.ts b/src/commands/ai/models/call.ts index 6bcc15e..9b066af 100644 --- a/src/commands/ai/models/call.ts +++ b/src/commands/ai/models/call.ts @@ -2,9 +2,9 @@ import color from '@heroku-cli/color' import {flags} from '@heroku-cli/command' import {Args, ux} from '@oclif/core' import fs from 'node:fs' -// import path from 'node:path' -import {ChatCompletionResponse, ModelList} from '../../../lib/ai/types' +import {ChatCompletionResponse, ImageResponse, ModelList} from '../../../lib/ai/types' import Command from '../../../lib/base' +import {openUrl} from '../../../lib/open-url' export default class Call extends Command { static args = { @@ -16,17 +16,18 @@ export default class Call extends Command { static description = 'make an inference request to a specific AI model resource' static examples = [ - 'heroku ai:models:call my_llm --prompt "What is the meaning of life?"', - 'heroku ai:models:call sdxl --prompt "Generate an image of a sunset" --opts \'{"quality": "hd"}\'', + 'heroku ai:models:call my_llm --app my-app --prompt "What is the meaning of life?"', + 'heroku ai:models:call sdxl --app my-app --prompt "Generate an image of a sunset" --opts \'{"quality":"hd"}\' -o sunset.png', ] static flags = { - app: flags.app({required: true}), + app: flags.app({required: false}), // interactive: flags.boolean({ // char: 'i', // description: 'Use interactive mode for conversation beyond the initial prompt (not available on all models)', // default: false, // }), + browser: flags.string({description: 'browser to open images with (example: "firefox", "safari")'}), json: flags.boolean({char: 'j', description: 'Output response as JSON'}), optfile: flags.string({ description: 'Additional options for model inference, provided as a JSON config file.', @@ -53,7 +54,7 @@ export default class Call extends Command { public async run(): Promise { const {args, flags} = await this.parse(Call) const {model_resource: modelResource} = args - const {app, json, optfile, opts, output, prompt} = flags + const {app, browser, json, optfile, opts, output, prompt} = flags // Initially, configure the default client to fetch the available model classes await this.configureHerokuAIClient() @@ -69,12 +70,15 @@ export default class Call extends Command { case 'Embedding': break - case 'Text-to-Image': + case 'Text-to-Image': { + const image = await this.generateImage(prompt, options) + await this.displayImageResult(image, output, browser, json) break + } case 'Text-to-Text': { const completion = await this.createChatCompletion(prompt, options) - this.displayChatCompletion(completion, output, json) + await this.displayChatCompletion(completion, output, json) break } @@ -146,13 +150,49 @@ export default class Call extends Command { return chatCompletionResponse } - private displayChatCompletion(completion: ChatCompletionResponse, output?: string, json = false) { - const content = json ? JSON.stringify(completion, null, 2) : completion.choices[0].message.content || '' + private async displayChatCompletion(completion: ChatCompletionResponse, output?: string, json = false) { + const content = completion.choices[0].message.content || '' if (output) { - fs.writeFileSync(output, content) + fs.writeFileSync(output, json ? JSON.stringify(completion, null, 2) : content) } else { json ? ux.styledJSON(completion) : ux.log(content) } } + + private async generateImage(prompt: string, options = {}) { + const {body: imageResponse} = await this.herokuAI.post('/v1/images/generations', { + body: { + ...options, + model: this.apiModelId, + prompt, + }, + headers: {authorization: `Bearer ${this.apiKey}`}, + }) + + return imageResponse + } + + private async displayImageResult(image: ImageResponse, output?: string, browser?: string, json = false) { + if (image.data[0].b64_json) { + if (output) { + const content = json ? JSON.stringify(image, null, 2) : Buffer.from(image.data[0].b64_json, 'base64') + fs.writeFileSync(output, content) + } else + json ? ux.styledJSON(image) : process.stdout.write(image.data[0].b64_json) + return + } + + if (image.data[0].url) { + if (output) + fs.writeFileSync(output, json ? JSON.stringify(image, null, 2) : image.data[0].url) + else if (json) + ux.styledJSON(image) + else + await openUrl(image.data[0].url, browser, 'view the image') + return + } + + ux.error('Unexpected response format') + } } diff --git a/src/lib/ai/types.ts b/src/lib/ai/types.ts index 56409fb..3935bce 100644 --- a/src/lib/ai/types.ts +++ b/src/lib/ai/types.ts @@ -150,3 +150,25 @@ export type ChatCompletionResponse = { } | null } } + +/** + * Image schema + */ +export type Image = { + /** The base64-encoded JSON of the generated image, if 'response_format' is 'b64_json' */ + readonly b64_json?: string | null + /** The prompt that was used to generate the image, if there was any revision to the prompt */ + readonly revised_prompt: string + /** The URL of the generated image, if 'response_format' is 'url' (default) */ + readonly url?: string | null +} + +/** + * Image response schema. + */ +export type ImageResponse = { + /** The Unix timestamp (in seconds) of when the image was generated */ + readonly created: number + /** A list of images */ + readonly data: Array +} diff --git a/src/lib/open-url.ts b/src/lib/open-url.ts new file mode 100644 index 0000000..27001a3 --- /dev/null +++ b/src/lib/open-url.ts @@ -0,0 +1,36 @@ +import color from '@heroku-cli/color' +import {ux} from '@oclif/core' +import {CLIError} from '@oclif/core/lib/errors' +import open from 'open' + +export const urlOpener: (...args: Parameters) => ReturnType = open + +export async function openUrl(url: string, browser?: string, action?: string) { + let browserErrorShown = false + const showBrowserError = (browser?: string) => { + if (browserErrorShown) return + + ux.warn(`Unable to open ${browser ? browser : 'your default'} browser. Please visit ${color.cyan(url)}${action ? ` to ${action}` : ''}.`) + browserErrorShown = true + } + + ux.log(`Opening ${color.cyan(url)} in ${browser ? browser : 'your default'} browser…`) + + try { + await ux.anykey( + `Press any key to open up the browser${action ? ` to ${action}` : ''}, or ${color.yellow('q')} to exit` + ) + } catch (error) { + const {message, oclif} = error as CLIError + ux.error(message, {exit: oclif?.exit || 1}) + } + + const cp = await urlOpener(url, {wait: false, ...(browser ? {app: {name: browser}} : {})}) + cp.on('error', (err: Error) => { + ux.warn(err) + showBrowserError(browser) + }) + cp.on('close', (code: number) => { + if (code !== 0) showBrowserError(browser) + }) +} diff --git a/test/commands/ai/docs.test.ts b/test/commands/ai/docs.test.ts index e420589..76ade26 100644 --- a/test/commands/ai/docs.test.ts +++ b/test/commands/ai/docs.test.ts @@ -1,21 +1,18 @@ -import {ux} from '@oclif/core' import {expect} from 'chai' -import childProcess from 'node:child_process' import sinon, {SinonSandbox, SinonStub} from 'sinon' -import {stderr, stdout} from 'stdout-stderr' import Cmd from '../../../src/commands/ai/docs' -import stripAnsi from '../../helpers/strip-ansi' import {runCommand} from '../../run-command' +import * as openUrl from '../../../src/lib/open-url' describe('ai:docs', function () { const {env} = process let sandbox: SinonSandbox - let urlOpener: SinonStub - let spawnMock: () => any + let openUrlStub: SinonStub beforeEach(function () { process.env = {} sandbox = sinon.createSandbox() + openUrlStub = sandbox.stub(openUrl, 'openUrl').onFirstCall().resolves() }) afterEach(function () { @@ -23,117 +20,33 @@ describe('ai:docs', function () { sandbox.restore() }) - context('when the user accepts the prompt to open the browser', function () { - beforeEach(function () { - sandbox.stub(ux, 'anykey').onFirstCall().resolves() - }) - - describe('attempting to open the browser', function () { - beforeEach(function () { - urlOpener = sandbox.stub(Cmd, 'urlOpener').onFirstCall().resolves({ - on: (_: string, _cb: ErrorCallback) => {}, - } as unknown as childProcess.ChildProcess) - }) - - context('without --browser option', function () { - it('shows the URL that will be opened for in the default browser', async function () { - await runCommand(Cmd) - - expect(stdout.output).to.contain(`Opening ${Cmd.defaultUrl} in your default browser…`) - }) - - it('attempts to open the default browser to the Dev Center AI article', async function () { - await runCommand(Cmd) - - expect(urlOpener.calledWith(Cmd.defaultUrl, {wait: false})).to.equal(true) - }) - }) - - context('with --browser option', function () { - it('shows the URL that will be opened in the specified browser', async function () { - await runCommand(Cmd, [ - '--browser=firefox', - ]) - - expect(stdout.output).to.contain(`Opening ${Cmd.defaultUrl} in firefox browser…`) - }) - - it('attempts to open the specified browser to the Dev Center AI article', async function () { - await runCommand(Cmd, [ - '--browser=firefox', - ]) - - expect(urlOpener.calledWith(Cmd.defaultUrl, {wait: false, app: {name: 'firefox'}})).to.equal(true) - }) - }) - - it('respects HEROKU_AI_DOCS_URL', async function () { - const customUrl = 'https://devcenter.heroku.com/articles/custom-article-url' - - process.env = { - HEROKU_AI_DOCS_URL: customUrl, - } + context('without --browser option', function () { + it('attempts to open the default browser to the Dev Center AI article', async function () { + await runCommand(Cmd) - await runCommand(Cmd) - - expect(urlOpener.calledWith(customUrl, {wait: false})).to.equal(true) - }) - }) - - context('when there’s an error opening the browser', function () { - beforeEach(function () { - spawnMock = sandbox.stub().returns({ - on: (event: string, cb: CallableFunction) => { - if (event === 'error') cb(new Error('error')) - }, unref: () => {}, - }) - }) - - it('shows a warning', async function () { - const spawnStub = sandbox.stub(childProcess, 'spawn').callsFake(spawnMock) - - await runCommand(Cmd) - - expect(spawnStub.calledOnce).to.be.true - expect(stripAnsi(stderr.output)).to.contain('Error: error') - expect(stripAnsi(stderr.output)).to.contain('Warning: Unable to open your default browser.') - expect(stripAnsi(stderr.output)).to.contain(Cmd.defaultUrl) - }) + expect(openUrlStub.calledWith(Cmd.defaultUrl, undefined, 'view the documentation')).to.be.true }) + }) - context('when the browser closes with a non-zero exit status', function () { - beforeEach(function () { - spawnMock = sandbox.stub().returns({ - on: (event: string, cb: CallableFunction) => { - if (event === 'close') cb(1) - }, unref: () => {}, - }) - }) - - it('shows a warning', async function () { - const spawnStub = sandbox.stub(childProcess, 'spawn').callsFake(spawnMock) - - await runCommand(Cmd) + context('with --browser option', function () { + it('attempts to open the specified browser to the Dev Center AI article', async function () { + await runCommand(Cmd, [ + '--browser=firefox', + ]) - expect(spawnStub.calledOnce).to.be.true - expect(stripAnsi(stderr.output)).to.contain('Warning: Unable to open your default browser.') - expect(stripAnsi(stderr.output)).to.contain(Cmd.defaultUrl) - }) + expect(openUrlStub.calledWith(Cmd.defaultUrl, 'firefox', 'view the documentation')).to.be.true }) }) - context('when the user rejects the prompt to open the browser', function () { - beforeEach(function () { - urlOpener = sandbox.stub(Cmd, 'urlOpener') - sandbox.stub(ux, 'anykey').onFirstCall().rejects(new Error('quit')) - }) + it('respects HEROKU_AI_DOCS_URL', async function () { + const customUrl = 'https://devcenter.heroku.com/articles/custom-article-url' - it('doesn’t attempt to open the browser', async function () { - try { - await runCommand(Cmd) - } catch {} + process.env = { + HEROKU_AI_DOCS_URL: customUrl, + } - expect(urlOpener.notCalled).to.equal(true) - }) + await runCommand(Cmd) + + expect(openUrlStub.calledWith(customUrl, undefined, 'view the documentation')).to.be.true }) }) diff --git a/test/commands/ai/models/call.test.ts b/test/commands/ai/models/call.test.ts index 0c0e2f7..b93f17f 100644 --- a/test/commands/ai/models/call.test.ts +++ b/test/commands/ai/models/call.test.ts @@ -4,9 +4,18 @@ import {expect} from 'chai' import nock from 'nock' import sinon from 'sinon' import Cmd from '../../../../src/commands/ai/models/call' +import * as openUrl from '../../../../src/lib/open-url' import stripAnsi from '../../../helpers/strip-ansi' import {runCommand} from '../../../run-command' -import {addon3, addon3Attachment1, availableModels, chatCompletionResponse} from '../../../helpers/fixtures' +import { + addon3, addon3Attachment1, + addon5, addon5Attachment1, + availableModels, + chatCompletionResponse, + mockedImageBase64, mockedImageContent, mockedImageResponseBase64, + mockedImageResponseUrl, + mockedImageUrl, +} from '../../../helpers/fixtures' import heredoc from 'tsheredoc' describe('ai:models:call', function () { @@ -48,7 +57,7 @@ describe('ai:models:call', function () { }) }) - context('without any optional flags', function () { + context('without --json or --output options', function () { it('sends the prompt to the service and displays the response content', async function () { const prompt = 'Hello, who are you?' inferenceApi = nock('https://inference-eu.heroku.com', { @@ -257,4 +266,223 @@ describe('ai:models:call', function () { }) }) }) + + context('when targeting a diffusion (Text-to-Image) model resource', function () { + beforeEach(async function () { + api.post('/actions/addons/resolve', {addon: addon5Attachment1.name, app: addon5Attachment1.app?.name}) + .reply(200, [addon5]) + .post('/actions/addon-attachments/resolve', {addon_attachment: addon5Attachment1.name, app: addon5Attachment1.app?.name}) + .reply(200, [addon5Attachment1]) + .get(`/apps/${addon5Attachment1.app?.id}/config-vars`) + .reply(200, { + DIFFUSION_KEY: 's3cr3t_k3y', + DIFFUSION_MODEL_ID: 'stable-diffusion-xl', + DIFFUSION_URL: 'inference-eu.heroku.com', + }) + }) + + context('without --json or --output options, for Base64 response format', function () { + it('sends the prompt to the service and shows the Base64-encoded content of the file', async function () { + const prompt = 'Generate a mocked image' + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + response_format: 'b64_json', + }).reply(200, mockedImageResponseBase64) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--opts={"response_format":"b64_json"}', + ]) + + expect(stdout.output).to.eq(mockedImageBase64) + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('with --json flag, but no --output, for Base64 response format', function () { + it('sends the prompt to the service and shows the JSON response', async function () { + const prompt = 'Generate a mocked image' + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + response_format: 'b64_json', + }).reply(200, mockedImageResponseBase64) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--opts={"response_format":"b64_json"}', + '--json', + ]) + + expect(JSON.parse(stdout.output)).to.deep.equal(mockedImageResponseBase64) + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('with --output option, but no --json, for Base64 response format', function () { + it('sends the prompt to the service, decodes the Base64 content and writes it to the indicated file', async function () { + const prompt = 'Generate a mocked image' + const writeFileSyncMock = sandbox.stub(fs, 'writeFileSync') + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + response_format: 'b64_json', + }).reply(200, mockedImageResponseBase64) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--opts={"response_format":"b64_json"}', + '--output=output-image.png', + ]) + + expect(writeFileSyncMock.calledWith( + 'output-image.png', + Buffer.from(mockedImageContent), + )).to.be.true + expect(stdout.output).to.eq('') + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('with --output and --json options, for Base64 response format', function () { + it('sends the prompt to the service and writes full JSON response to the indicated file', async function () { + const prompt = 'Generate a mocked image' + const writeFileSyncMock = sandbox.stub(fs, 'writeFileSync') + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + response_format: 'b64_json', + }).reply(200, mockedImageResponseBase64) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--opts={"response_format":"b64_json"}', + '--output=image-response.json', + '--json', + ]) + + expect(writeFileSyncMock.calledWith( + 'image-response.json', + JSON.stringify(mockedImageResponseBase64, null, 2), + )).to.be.true + expect(stdout.output).to.eq('') + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('without --json or --output options, for URL response format', function () { + it('sends the prompt to the service and attempts to open the URL on the default browser', async function () { + const openUrlStub = sandbox.stub(openUrl, 'openUrl').onFirstCall().resolves() + const prompt = 'Generate a mocked image' + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + }).reply(200, mockedImageResponseUrl) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + ]) + + expect(openUrlStub.calledWith(mockedImageUrl, undefined, 'view the image')).to.be.true + }) + }) + + context('with --json flag, but no --output, for URL response format', function () { + it('sends the prompt to the service and shows the JSON response', async function () { + const prompt = 'Generate a mocked image' + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + }).reply(200, mockedImageResponseUrl) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--json', + ]) + + expect(JSON.parse(stdout.output)).to.deep.equal(mockedImageResponseUrl) + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('with --output option, but no --json, for URL response format', function () { + it('sends the prompt to the service and writes the URL to the indicated file', async function () { + const prompt = 'Generate a mocked image' + const writeFileSyncMock = sandbox.stub(fs, 'writeFileSync') + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + }).reply(200, mockedImageResponseUrl) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--output=image-url.txt', + ]) + + expect(writeFileSyncMock.calledWith( + 'image-url.txt', + mockedImageUrl, + )).to.be.true + expect(stdout.output).to.eq('') + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + + context('with --output and --json options, for URL response format', function () { + it('sends the prompt to the service and writes full JSON response to the indicated file', async function () { + const prompt = 'Generate a mocked image' + const writeFileSyncMock = sandbox.stub(fs, 'writeFileSync') + inferenceApi = nock('https://inference-eu.heroku.com', { + reqheaders: {authorization: 'Bearer s3cr3t_k3y'}, + }).post('/v1/images/generations', { + model: 'stable-diffusion-xl', + prompt, + }).reply(200, mockedImageResponseUrl) + + await runCommand(Cmd, [ + 'DIFFUSION', + '--app=app2', + `--prompt=${prompt}`, + '--output=image-response.json', + '--json', + ]) + + expect(writeFileSyncMock.calledWith( + 'image-response.json', + JSON.stringify(mockedImageResponseUrl, null, 2), + )).to.be.true + expect(stdout.output).to.eq('') + expect(stripAnsi(stderr.output)).to.eq('') + }) + }) + }) }) diff --git a/test/helpers/fixtures.ts b/test/helpers/fixtures.ts index ad5e4ac..c76d68a 100644 --- a/test/helpers/fixtures.ts +++ b/test/helpers/fixtures.ts @@ -1,5 +1,5 @@ import * as Heroku from '@heroku-cli/schema' -import {ChatCompletionResponse} from '../../src/lib/ai/types' +import {ChatCompletionResponse, ImageResponse} from '../../src/lib/ai/types' export const availableModels = [ { @@ -241,7 +241,7 @@ export const addon1ProvisionedWithAttachmentName: Heroku.AddOn = { export const chatCompletionResponse: ChatCompletionResponse = { id: 'chatcmpl-17f8f365f941de720ad38', object: 'chat.completion', - created: 1727398076, + created: 1234567890, model: 'claude-3-sonnet', system_fingerprint: 'heroku-inf-zzuqrd', choices: [ @@ -262,3 +262,56 @@ export const chatCompletionResponse: ChatCompletionResponse = { }, } +export const addon5: Heroku.AddOn = { + addon_service: { + id: '4b46be3f-d0e6-4b3f-b616-0a857115d71d', + name: 'inference', + }, + app: { + id: 'd0256f69-a6ea-45ad-93e5-3911eac0d216', + name: 'app2', + }, + id: 'c0addaa4-d2e2-4da4-bf93-c522af6790f9', + name: 'inference-colorful-79696', + plan: { + id: 'de948fb0-48c4-4f47-912d-745817a80f05', + name: 'inference:stable-diffusion-xl', + }, +} + +export const addon5Attachment1: Heroku.AddOnAttachment = { + addon: { + id: 'c0addaa4-d2e2-4da4-bf93-c522af6790f9', + name: 'inference-colorful-79696', + app: { + id: 'd0256f69-a6ea-45ad-93e5-3911eac0d216', + name: 'app2', + }, + }, + app: { + id: 'd0256f69-a6ea-45ad-93e5-3911eac0d216', + name: 'app2', + }, + id: '87f6f66f-8ad3-4787-b895-bc79c2641342', + name: 'DIFFUSION', +} + +export const mockedImageContent = 'Let’s pretend this is an image' +export const mockedImageBase64 = 'TGV04oCZcyBwcmV0ZW5kIHRoaXMgaXMgYW4gaW1hZ2U=' +export const mockedImageUrl = 'https://example.com/image.png' + +export const mockedImageResponseBase64: ImageResponse = { + created: 1234567890, + data: [{ + b64_json: mockedImageBase64, + revised_prompt: '', + }], +} + +export const mockedImageResponseUrl: ImageResponse = { + created: 1234567890, + data: [{ + url: mockedImageUrl, + revised_prompt: '', + }], +} diff --git a/test/lib/open-url.test.ts b/test/lib/open-url.test.ts new file mode 100644 index 0000000..82b102c --- /dev/null +++ b/test/lib/open-url.test.ts @@ -0,0 +1,159 @@ +import {ux} from '@oclif/core' +import {expect} from 'chai' +import childProcess from 'node:child_process' +import sinon, {SinonSandbox, SinonStub} from 'sinon' +import {stderr, stdout} from 'stdout-stderr' +import stripAnsi from '../helpers/strip-ansi' +import * as openUrl from '../../src/lib/open-url' + +const stdOutputMockStart = () => { + stderr.start() + stdout.start() +} + +const stdOutputMockStop = () => { + stderr.stop() + stdout.stop() +} + +describe('open-url', function () { + const {env} = process + const url = 'https://example.com' + let sandbox: SinonSandbox + let urlOpenerStub: SinonStub + let anyKeyStub: SinonStub + let spawnMock: () => any + + beforeEach(function () { + process.env = {} + sandbox = sinon.createSandbox() + }) + + afterEach(function () { + sandbox.restore() + process.env = env + }) + + context('when the user accepts the prompt to open the browser', function () { + beforeEach(function () { + anyKeyStub = sandbox.stub(ux, 'anykey').onFirstCall().resolves() + }) + + describe('attempting to open the browser', function () { + beforeEach(function () { + urlOpenerStub = sandbox.stub(openUrl, 'urlOpener').onFirstCall().resolves({ + on: (_: string, _cb: ErrorCallback) => {}, + } as unknown as childProcess.ChildProcess) + }) + + context('without browser or action arguments', function () { + it('shows the URL that will be opened for in the default browser', async function () { + stdOutputMockStart() + await openUrl.openUrl(url) + stdOutputMockStop() + + expect(stdout.output).to.contain(`Opening ${url} in your default browser…`) + }) + + it('attempts to open the default browser to the url argument', async function () { + stdOutputMockStart() + await openUrl.openUrl(url) + stdOutputMockStop() + + expect(urlOpenerStub.calledWith(url, {wait: false})).to.equal(true) + }) + }) + + context('with browser argument', function () { + it('shows the URL that will be opened in the specified browser', async function () { + stdOutputMockStart() + await openUrl.openUrl(url, 'firefox') + stdOutputMockStop() + + expect(stdout.output).to.contain(`Opening ${url} in firefox browser…`) + }) + + it('attempts to open the specified browser to the url argument', async function () { + stdOutputMockStart() + await openUrl.openUrl(url, 'firefox') + stdOutputMockStop() + + expect(urlOpenerStub.calledWith(url, {wait: false, app: {name: 'firefox'}})).to.equal(true) + }) + }) + + context('with action argument', function () { + it('shows the action to be performed', async function () { + stdOutputMockStart() + await openUrl.openUrl(url, undefined, 'view something') + stdOutputMockStop() + + expect(anyKeyStub.calledWithMatch(/to view something/)).to.be.true + }) + }) + }) + + context('when there’s an error opening the browser', function () { + beforeEach(function () { + spawnMock = sandbox.stub().returns({ + on: (event: string, cb: CallableFunction) => { + if (event === 'error') cb(new Error('error')) + }, unref: () => {}, + }) + }) + + it('shows a warning', async function () { + const spawnStub = sandbox.stub(childProcess, 'spawn').callsFake(spawnMock) + + stdOutputMockStart() + await openUrl.openUrl(url) + stdOutputMockStop() + + expect(spawnStub.calledOnce).to.be.true + expect(stripAnsi(stderr.output)).to.contain('Error: error') + expect(stripAnsi(stderr.output)).to.contain('Warning: Unable to open your default browser.') + expect(stripAnsi(stderr.output)).to.contain(url) + }) + }) + + context('when the browser closes with a non-zero exit status', function () { + beforeEach(function () { + spawnMock = sandbox.stub().returns({ + on: (event: string, cb: CallableFunction) => { + if (event === 'close') cb(1) + }, unref: () => {}, + }) + }) + + it('shows a warning', async function () { + const spawnStub = sandbox.stub(childProcess, 'spawn').callsFake(spawnMock) + + stdOutputMockStart() + await openUrl.openUrl(url) + stdOutputMockStop() + + expect(spawnStub.calledOnce).to.be.true + expect(stripAnsi(stderr.output)).to.contain('Warning: Unable to open your default browser.') + expect(stripAnsi(stderr.output)).to.contain(url) + }) + }) + }) + + context('when the user rejects the prompt to open the browser', function () { + beforeEach(function () { + urlOpenerStub = sandbox.stub(openUrl, 'urlOpener') + sandbox.stub(ux, 'anykey').onFirstCall().rejects(new Error('quit')) + }) + + it('doesn’t attempt to open the browser', async function () { + try { + stdOutputMockStart() + await openUrl.openUrl(url) + } catch { + stdOutputMockStop() + } + + expect(urlOpenerStub.notCalled).to.equal(true) + }) + }) +})