diff --git a/src/commands/config.ts b/src/commands/config.ts index 2ae38151..910695e0 100644 --- a/src/commands/config.ts +++ b/src/commands/config.ts @@ -37,7 +37,7 @@ export default command( } if (mode === 'get') { - const config = await getConfig(); + const config = await getConfig(undefined, { allowMissingKey: true }); for (const key of keyValues) { if (hasOwn(config, key)) { console.log(`${key}=${config[key as keyof typeof config]}`); diff --git a/src/helpers/completion.ts b/src/helpers/completion.ts index 33171961..2535c9cb 100644 --- a/src/helpers/completion.ts +++ b/src/helpers/completion.ts @@ -18,7 +18,7 @@ import readline from 'readline'; const explainInSecondRequest = true; -function getOpenAi(key: string, apiEndpoint: string) { +function getOpenAi(key: string | undefined, apiEndpoint: string) { const openAi = new OpenAIApi( new Configuration({ apiKey: key, basePath: apiEndpoint }) ); @@ -320,11 +320,16 @@ function getRevisionPrompt(prompt: string, code: string) { } export async function getModels( - key: string, + key: string | undefined, apiEndpoint: string ): Promise { const openAi = getOpenAi(key, apiEndpoint); const response = await openAi.listModels(); - return response.data.data.filter((model) => model.object === 'model'); + return response.data.data.filter( + (model) => + model.object === 'model' && + typeof model.id === 'string' && + model.id.trim().length > 0 + ); } diff --git a/src/helpers/config.ts b/src/helpers/config.ts index cf3549af..a6b5b23c 100644 --- a/src/helpers/config.ts +++ b/src/helpers/config.ts @@ -66,6 +66,10 @@ type ValidConfig = { [Key in ConfigKeys]: ReturnType<(typeof configParsers)[Key]>; }; +type GetConfigOptions = { + allowMissingKey?: boolean; +}; + const configPath = path.join(os.homedir(), '.ai-shell'); const fileExists = (filePath: string) => @@ -85,7 +89,8 @@ const readConfigFile = async (): Promise => { }; export const getConfig = async ( - cliConfig?: RawConfig + cliConfig?: RawConfig, + options: GetConfigOptions = {} ): Promise => { const config = await readConfigFile(); const parsedConfig: Record = {}; @@ -93,6 +98,9 @@ export const getConfig = async ( for (const key of Object.keys(configParsers) as ConfigKeys[]) { const parser = configParsers[key]; const value = cliConfig?.[key] ?? config[key]; + if (key === 'OPENAI_KEY' && options.allowMissingKey && !value) { + continue; + } parsedConfig[key] = parser(value); } @@ -114,9 +122,23 @@ export const setConfigs = async (keyValues: [key: string, value: string][]) => { await fs.writeFile(configPath, ini.stringify(config), 'utf8'); }; +const customModelValue = '__custom_model__'; + +export function getModelSelectionOptions(models: Pick[]) { + return [ + ...models + .filter((model) => typeof model.id === 'string' && model.id.trim()) + .map((model) => ({ value: model.id, label: model.id })), + { + label: i18n.t('Enter the model you want to use'), + value: customModelValue, + }, + ]; +} + export const showConfigUI = async () => { try { - const config = await getConfig(); + const config = await getConfig(undefined, { allowMissingKey: true }); const choice = (await p.select({ message: i18n.t('Set config') + ':', options: [ @@ -189,16 +211,42 @@ export const showConfigUI = async () => { await setConfigs([['SILENT_MODE', silentMode ? 'true' : 'false']]); } else if (choice === 'MODEL') { 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; + await getConfig(undefined, { allowMissingKey: true }); + let models: Model[] = []; + if (key) { + try { + models = await getModels(key, apiEndpoint); + } catch { + // A custom endpoint may not implement /models. Offer manual entry. + } + } + + let model = models.length + ? ((await p.select({ + message: 'Pick a model.', + options: getModelSelectionOptions(models), + })) as string) + : ((await p.text({ + message: i18n.t('Enter the model you want to use'), + validate: (value) => { + if (!value.trim()) { + return i18n.t('Please enter a prompt.'); + } + }, + })) as string); if (p.isCancel(model)) return; + if (model === customModelValue) { + model = (await p.text({ + message: i18n.t('Enter the model you want to use'), + validate: (value) => { + if (!value.trim()) { + return i18n.t('Please enter a prompt.'); + } + }, + })) as string; + if (p.isCancel(model)) return; + } await setConfigs([['MODEL', model]]); } else if (choice === 'LANGUAGE') { const language = (await p.select({ diff --git a/tests/config-model-fallback.test.ts b/tests/config-model-fallback.test.ts new file mode 100644 index 00000000..28b09e0a --- /dev/null +++ b/tests/config-model-fallback.test.ts @@ -0,0 +1,17 @@ +import assert from 'node:assert/strict'; +import { getModelSelectionOptions } from '../src/helpers/config'; + +const emptyOptions = getModelSelectionOptions([]); +assert.equal(emptyOptions.length, 1); +assert.equal(emptyOptions[0].value, '__custom_model__'); + +const options = getModelSelectionOptions([ + { id: 'gpt-4o', object: 'model' }, + { id: '', object: 'model' }, +]); +assert.deepEqual(options, [ + { value: 'gpt-4o', label: 'gpt-4o' }, + { label: 'Enter the model you want to use', value: '__custom_model__' }, +]); + +console.log('config model fallback checks passed');