Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(plugin-ai): Add 'ai:models:info' command #11

Merged
merged 22 commits into from
Sep 28, 2024
Merged
Show file tree
Hide file tree
Changes from 19 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ USAGE
* [`heroku ai:models:create MODEL_NAME`](#heroku-aimodelscreate-model_name)
* [`heroku ai:models:destroy MODELRESOURCE`](#heroku-aimodelsdestroy-modelresource)
* [`heroku ai:models:detach MODEL_RESOURCE`](#heroku-aimodelsdetach-model_resource)
* [`heroku ai:models:info [MODELRESOURCE]`](#heroku-aimodelsinfo-modelresource)
* [`heroku ai:models:list`](#heroku-aimodelslist)

## `heroku ai:docs`
Expand Down Expand Up @@ -175,6 +176,32 @@ EXAMPLES

_See code: [dist/commands/ai/models/detach.ts](https://github.com/heroku/heroku-cli-plugin-integration/blob/v0.0.0/dist/commands/ai/models/detach.ts)_

## `heroku ai:models:info [MODELRESOURCE]`

get the current status of all the AI model resources attached to your app or a specific resource

```
USAGE
$ heroku ai:models:info [MODELRESOURCE] -a <value> [-r <value>]

ARGUMENTS
MODELRESOURCE The resource ID or alias of the model resource to check.

FLAGS
-a, --app=<value> (required) app to run command against
-r, --remote=<value> git remote of app to use

DESCRIPTION
get the current status of all the AI model resources attached to your app or a specific resource

EXAMPLES
$ heroku ai:models:info claude-3-5-sonnet-acute-04281 --app example-app

$ heroku ai:models:info --app example-app
```

_See code: [dist/commands/ai/models/info.ts](https://github.com/heroku/heroku-cli-plugin-integration/blob/v0.0.0/dist/commands/ai/models/info.ts)_

## `heroku ai:models:list`

list available AI models to provision access to
Expand Down
1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
"lint": "eslint . --ext .ts --config .eslintrc.json",
"posttest": "yarn lint",
"test": "nyc mocha --forbid-only",
"test:local": "nyc mocha",
"version": "oclif readme && git add README.md"
}
}
104 changes: 104 additions & 0 deletions src/commands/ai/models/info.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
import color from '@heroku-cli/color'
import {flags} from '@heroku-cli/command'
import {Args, ux} from '@oclif/core'
import Command from '../../../lib/base'
import {ModelResource} from '../../../lib/ai/types'
import appAddons from '../../../lib/ai/models/app_addons'
import * as Heroku from '@heroku-cli/schema'

export default class Info extends Command {
static description = 'get the current status of all the AI model resources attached to your app or a specific resource'
static examples = [
'heroku ai:models:info claude-3-5-sonnet-acute-04281 --app example-app',
'heroku ai:models:info --app example-app',
]

static flags = {
app: flags.app({required: true}),
remote: flags.remote(),
}

static args = {
modelResource: Args.string({description: 'The resource ID or alias of the model resource to check.'}),
}

public async run(): Promise<any> {
const {args, flags} = await this.parse(Info)
const {app} = flags
const {modelResource} = args
const synthesizedModels: Array<ModelResource> = []
let listOfProvisionedModels: Array<ModelResource> = []

const modelInfo = async () => {
const modelInfoResponse = await this.herokuAI.get<ModelResource>(`/models/${this.apiModelId}`, {
headers: {authorization: `Bearer ${this.apiKey}`},
})
.catch(error => {
if (error.statusCode === 404) {
ux.warn(`We can’t find a model resource called ${color.yellow(modelResource)}.\nRun ${color.cmd('heroku ai:models:info -a <app>')} to see a list of model resources.`)
} else {
throw error
}
})

return modelInfoResponse
}

const getModelDetails = async (collectedModels: Array<Heroku.AddOn> | string) => {
if (typeof collectedModels === 'string') {
const modelResource = collectedModels
await this.configureHerokuAIClient(modelResource, app)

const {body: currentModelResource} = await modelInfo() || {body: null}
synthesizedModels.push(currentModelResource!)
} else {
for (const addonModel of collectedModels) {
await this.configureHerokuAIClient(addonModel.modelResource, app)

const {body: currentModelResource} = await modelInfo() || {body: null}
synthesizedModels.push(currentModelResource!)
}
}

return synthesizedModels
}

if (modelResource) {
listOfProvisionedModels = await getModelDetails(modelResource)
} else {
const provisionedModelsInfo: Record<string, string | undefined>[] = []
const inferenceRegex = /inference/
const addonsResponse = await appAddons(this.config, app)

for (const addonInfo of addonsResponse as Array<Heroku.AddOn>) {
const addonType = addonInfo.addon_service?.name || ''
const isModelAddon = inferenceRegex.test(addonType)

if (isModelAddon) {
provisionedModelsInfo.push({
addonName: addonInfo.addon_service?.name,
modelResource: addonInfo.name,
modelId: addonInfo.addon_service?.id,
})
}
}

listOfProvisionedModels = await getModelDetails(provisionedModelsInfo)
}

this.displayModelResource(listOfProvisionedModels)
}

displayModelResource(modelResources: ModelResource[]) {
for (const modelResource of modelResources) {
ux.log()
ux.styledObject({
'Base Model ID': modelResource.plan,
Ready: modelResource.ready,
'Tokens In': modelResource.tokens_in,
'Tokens Out': modelResource.tokens_out,
'Avg Performance': modelResource.avg_performance,
})
}
}
}
17 changes: 17 additions & 0 deletions src/lib/ai/models/app_addons.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import {Config} from '@oclif/core'
import {APIClient} from '@heroku-cli/command'
import * as Heroku from '@heroku-cli/schema'

export default async function (config: Config, app: string) {
const herokuClient = new APIClient(config)

const {body: response} = await herokuClient.get<Heroku.AddOn>(`/apps/${app}/addons`, {
headers: {'Accept-Expansion': 'plan'},
}).catch(error => {
console.log('ERROR MESSAGE:', error.message)
const error_ = error.body && error.body.message ? new Error(`The add-on was unable to be destroyed: ${error.body.message}.`) : new Error(`The add-on was unable to be destroyed: ${error}.`)
throw error_
})

return response
}
1 change: 1 addition & 0 deletions src/lib/ai/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ export type ModelInfo = {
*/
export type ModelResource = {
plan: ModelName
ready: string
created: string
tokens_in: string
tokens_out?: string
Expand Down
3 changes: 2 additions & 1 deletion src/lib/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -186,8 +186,9 @@ export default abstract class extends Command {
}

// 5. If we resolved for an add-on, check that it's a Managed Inference add-on or throw a NotFound error.
if (resolvedAddon && resolvedAddon.addon_service.name !== this.addonServiceSlug)
if (resolvedAddon && resolvedAddon.addon_service.name !== this.addonServiceSlug) {
throw new NotFound(addonIdentifier, appIdentifier)
}

// 6. If we resolved for an add-on but not for an attachment yet, try to resolve the attachment
if (resolvedAddon && !resolvedAttachment) {
Expand Down
163 changes: 163 additions & 0 deletions test/commands/ai/models/info.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,163 @@
import {expect} from 'chai'
import {stdout, stderr} from 'stdout-stderr'
import Cmd from '../../../../src/commands/ai/models/info'
import {runCommand} from '../../../run-command'
import {modelResource, addon1Attachment1, addon1, mockAPIErrors} from '../../../helpers/fixtures'
import nock from 'nock'
import heredoc from 'tsheredoc'
import stripAnsi from '../../../helpers/strip-ansi'
import {CLIError} from '@oclif/core/lib/errors'

describe('ai:models:info', function () {
const {env} = process
let api: nock.Scope
let herokuAI: nock.Scope

context('when provisioned model name is provided and is found', function () {
beforeEach(function () {
process.env = {}
api = nock('https://api.heroku.com:443')
herokuAI = nock('https://inference.heroku.com')
})

afterEach(function () {
process.env = env
nock.cleanAll()
})

it('shows info for a model resource', async function () {
api
.post('/actions/addons/resolve',
{addon: addon1.name, app: addon1Attachment1.app?.name})
.reply(200, [addon1])
.get(`/addons/${addon1.id}/addon-attachments`)
.reply(200, [addon1Attachment1])
.get(`/apps/${addon1Attachment1.app?.id}/config-vars`)
.reply(200, {
INFERENCE_KEY: 's3cr3t_k3y',
INFERENCE_MODEL_ID: 'claude-3-haiku',
INFERENCE_URL: 'inference.heroku.com',
})
herokuAI
.get('/models/claude-3-haiku')
.reply(200, modelResource)

await runCommand(Cmd, [
'inference-regular-74659',
'--app',
'app1',
])

expect(stripAnsi(stdout.output)).to.contain(heredoc`
Avg Performance: latency 0.4sec, 28 tokens/sec
Base Model ID: claude-3-haiku
Ready: Yes
Tokens In: 0 tokens this period
Tokens Out: 0 tokens this period
`)

expect(stderr.output).to.eq('')
})
})

context('when provisioned model name is not provided', function () {
// eslint-disable-next-line mocha/no-setup-in-describe
const multipleAddons = Array.from({length: 2}).fill(addon1)

beforeEach(function () {
process.env = {}
api = nock('https://api.heroku.com:443')
})

afterEach(function () {
process.env = env
nock.cleanAll()
})

it('shows info for all model resources on specified app', async function () {
api
.post('/actions/addons/resolve',
{addon: addon1.name, app: addon1Attachment1.app?.name})
.reply(200, [addon1])
.get(`/addons/${addon1.id}/addon-attachments`)
.reply(200, [addon1Attachment1])
.get(`/apps/${addon1Attachment1.app?.id}/config-vars`)
.reply(200, {
INFERENCE_KEY: 's3cr3t_k3y',
INFERENCE_MODEL_ID: 'claude-3-haiku',
INFERENCE_URL: 'inference.heroku.com',
})
herokuAI
.get('/models/claude-3-haiku')
.reply(200, modelResource)
api
.get(`/apps/${addon1.app?.name}/addons`)
.reply(200, multipleAddons)
.post('/actions/addons/resolve',
{addon: addon1.name, app: addon1Attachment1.app?.name})
.reply(200, [addon1])
.get(`/addons/${addon1.id}/addon-attachments`)
.reply(200, [addon1Attachment1])
.get(`/apps/${addon1Attachment1.app?.id}/config-vars`)
.reply(200, {
INFERENCE_KEY: 's3cr3t_k3y',
INFERENCE_MODEL_ID: 'claude-3-haiku',
INFERENCE_URL: 'inference.heroku.com',
})
herokuAI
.get('/models/claude-3-haiku')
.reply(200, modelResource)

await runCommand(Cmd, [
'--app',
'app1',
])

expect(stdout.output).to.contain(heredoc`
Avg Performance: latency 0.4sec, 28 tokens/sec
Base Model ID: claude-3-haiku
Ready: Yes
Tokens In: 0 tokens this period
Tokens Out: 0 tokens this period

Avg Performance: latency 0.4sec, 28 tokens/sec
Base Model ID: claude-3-haiku
Ready: Yes
Tokens In: 0 tokens this period
Tokens Out: 0 tokens this period
`)
})
})

context('when provisioned model name is incorrectly inputted', function () {
const incorrectModelName = 'inference-regular-WRONG'

beforeEach(function () {
process.env = {}
api = nock('https://api.heroku.com:443')
})

afterEach(function () {
process.env = env
nock.cleanAll()
})

it('shows an error message', async function () {
api
.post('/actions/addons/resolve',
{addon: incorrectModelName, app: addon1Attachment1.app?.name})
.reply(404, mockAPIErrors.modelsInfoErrorResponse)

try {
await runCommand(Cmd, [
incorrectModelName,
'--app',
'app1',
])
} catch (error) {
const {message} = error as CLIError
expect(stripAnsi(message)).contains(mockAPIErrors.modelsInfoErrorResponse.message)
}
})
})
})
14 changes: 14 additions & 0 deletions test/helpers/fixtures.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import * as Heroku from '@heroku-cli/schema'
import {ModelResource} from '../../src/lib/ai/types'

export const availableModels = [
{
Expand Down Expand Up @@ -46,6 +47,19 @@ export const mockAPIErrors = {
id: 'error',
message: 'Example API Error',
},
modelsInfoErrorResponse: {
id: 'error',
message: 'Example API Error',
},
}

export const modelResource: ModelResource = {
plan: 'claude-3-haiku',
ready: 'Yes',
created: '2023-01-21T13:02:37.320+00.00',
tokens_in: '0 tokens this period',
tokens_out: '0 tokens this period',
avg_performance: 'latency 0.4sec, 28 tokens/sec',
}

export const addon1: Heroku.AddOn = {
Expand Down