From 7d3ccadc964ef864d69979f50e3ec0d0ad6f4a75 Mon Sep 17 00:00:00 2001 From: chen201724 <1227646458@qq.com> Date: Thu, 6 Jun 2024 18:35:10 +0800 Subject: [PATCH] Add support for user to select from currently available models --- CHANGELOG.md | 1 + src/helpers/completion.ts | 17 ++++++++++++++++- src/helpers/config.ts | 15 ++++++++++++--- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbd2308..b9551ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ ## 1.0.8 + - Fix a bug causing completions to not display in some cases - Allow hitting "q" or "esc" to stop the streaming output diff --git a/src/helpers/completion.ts b/src/helpers/completion.ts index 85c1591..7cd5c6f 100644 --- a/src/helpers/completion.ts +++ b/src/helpers/completion.ts @@ -1,4 +1,9 @@ -import { OpenAIApi, Configuration, ChatCompletionRequestMessage } from 'openai'; +import { + OpenAIApi, + Configuration, + ChatCompletionRequestMessage, + Model, +} from 'openai'; import dedent from 'dedent'; import { IncomingMessage } from 'http'; import { KnownError } from './error'; @@ -313,3 +318,13 @@ function getRevisionPrompt(prompt: string, code: string) { ${generationDetails} `; } + +export async function getModels( + key: string, + apiEndpoint: string +): Promise { + const openAi = getOpenAi(key, apiEndpoint); + const response = await openAi.listModels(); + + return response.data.data.filter((model) => model.object === 'model'); +} diff --git a/src/helpers/config.ts b/src/helpers/config.ts index a536270..8ac8c90 100644 --- a/src/helpers/config.ts +++ b/src/helpers/config.ts @@ -8,6 +8,8 @@ import { KnownError, handleCliError } from './error'; import * as p from '@clack/prompts'; import { red } from 'kolorist'; import i18n from './i18n'; +import { getModels } from './completion'; +import { Model } from 'openai'; const { hasOwnProperty } = Object.prototype; export const hasOwn = (object: unknown, key: PropertyKey) => @@ -186,9 +188,16 @@ export const showConfigUI = async () => { if (p.isCancel(silentMode)) return; await setConfigs([['SILENT_MODE', silentMode ? 'true' : 'false']]); } else if (choice === 'MODEL') { - const model = await p.text({ - message: i18n.t('Enter the model you want to use'), - }); + const { OPENAI_KEY: key, OPENAI_API_ENDPOINT: apiEndpoint } = + await getConfig(); + const models = await getModels(key, apiEndpoint); + const model = (await p.select({ + message: 'Pick a model.', + options: models.map((m: Model) => { + return { value: m.id, label: m.id }; + }), + })) as string; + if (p.isCancel(model)) return; await setConfigs([['MODEL', model]]); } else if (choice === 'LANGUAGE') {