Skip to content

Commit

Permalink
Add support for user to select from currently available models
Browse files Browse the repository at this point in the history
  • Loading branch information
chen201724 committed Jun 6, 2024
1 parent f1d1b1c commit 7d3ccad
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 4 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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

Expand Down
17 changes: 16 additions & 1 deletion src/helpers/completion.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -313,3 +318,13 @@ function getRevisionPrompt(prompt: string, code: string) {
${generationDetails}
`;
}

export async function getModels(
key: string,
apiEndpoint: string
): Promise<Model[]> {
const openAi = getOpenAi(key, apiEndpoint);
const response = await openAi.listModels();

return response.data.data.filter((model) => model.object === 'model');
}
15 changes: 12 additions & 3 deletions src/helpers/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) =>
Expand Down Expand Up @@ -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') {
Expand Down

0 comments on commit 7d3ccad

Please sign in to comment.