Skip to content

Commit

Permalink
feat: support custom keys for custom Open AI models
Browse files Browse the repository at this point in the history
The configuration for custom OpenAI models now allows specifying a
unique 'apiKey' for each model, or reusing the global OpenAI API key.

fixes #14288
  • Loading branch information
sdirix committed Oct 17, 2024
1 parent 78ce064 commit fc1b88e
Show file tree
Hide file tree
Showing 6 changed files with 58 additions and 20 deletions.
20 changes: 19 additions & 1 deletion packages/ai-openai/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,25 @@

The `@theia/ai-openai` integrates OpenAI's models with Theia AI.
The OpenAI API key and the models to use can be configured via preferences.
Alternatively the OpenAI API key can also be handed in via an environment variable.
Alternatively the OpenAI API key can also be handed in via the `OPENAI_API_KEY` variable.

### Custom models

The extension also supports OpenAI compatible models hosted on different end points.
You can configure the end points via the `ai-features.openAiCustom.customOpenAiModels` preference:

```ts
{
model: string
url: string
id?: string
apiKey?: string | true
}
```

- `model` and `url` are mandatory attributes, indicating the end point and model to use
- `id` is an optional attribute which is used in the UI to refer to this configuration
- `apiKey` is either the key to access the API served at the given URL or `true` to use the global OpenAI API key. If not given 'no-key' will be used.

## Additional Information

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ export class OpenAiFrontendApplicationContribution implements FrontendApplicatio

const modelsToRemove = oldModels.filter(model => !newModels.some(newModel => newModel.id === model.id));
const modelsToAddOrUpdate = newModels.filter(newModel => !oldModels.some(model =>
model.id === newModel.id && model.model === newModel.model && model.url === newModel.url));
model.id === newModel.id && model.model === newModel.model && model.url === newModel.url && model.apiKey === newModel.apiKey));

this.manager.removeLanguageModels(...modelsToRemove.map(model => model.id));
this.manager.createOrUpdateLanguageModels(...modelsToAddOrUpdate);
Expand All @@ -77,21 +77,23 @@ export class OpenAiFrontendApplicationContribution implements FrontendApplicatio
function createOpenAIModelDescription(modelId: string): OpenAiModelDescription {
return {
id: `openai/${modelId}`,
model: modelId
model: modelId,
apiKey: true
};
}

function createCustomModelDescriptionsFromPreferences(preferences: Partial<OpenAiModelDescription>[]): OpenAiModelDescription[] {
return preferences.reduce((acc, pref) => {
if (!pref.model || !pref.url) {
if (!pref.model || !pref.url || typeof pref.model !== 'string' || typeof pref.url !== 'string') {
return acc;
}
return [
...acc,
{
id: pref.id ?? pref.model,
id: pref.id && typeof pref.id === 'string' ? pref.id : pref.model,
model: pref.model,
url: pref.url
url: pref.url,
apiKey: typeof pref.apiKey === 'string' || pref.apiKey === true ? pref.apiKey : undefined
}
];
}, []);
Expand Down
13 changes: 11 additions & 2 deletions packages/ai-openai/src/browser/openai-preferences.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,12 @@ export const OpenAiPreferencesSchema: PreferenceSchema = {
type: 'array',
title: AI_CORE_PREFERENCES_TITLE,
markdownDescription: 'Integrate custom models compatible with the OpenAI API, for example via `vllm`. The required attributes are `model` and `url`.\
Optionally, you can provide a unique `id` to identify the custom model in the UI. If none is given `model` will be used as `id`.',
\n\
Optionally, you can\
\n\
- specify a unique `id` to identify the custom model in the UI. If none is given `model` will be used as `id`.\
\n\
- provide an `apiKey` to access the API served at the given url. Use `true` to indicate the use of the global OpenAI API key.',
default: [],
items: {
type: 'object',
Expand All @@ -59,7 +64,11 @@ export const OpenAiPreferencesSchema: PreferenceSchema = {
id: {
type: 'string',
title: 'A unique identifier which is used in the UI to identify the custom model',
}
},
apiKey: {
type: ['string', 'boolean'],
title: 'Either the key to access the API served at the given url or `true` to use the global OpenAI API key',
},
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,10 @@ export interface OpenAiModelDescription {
* The OpenAI API compatible endpoint where the model is hosted. If not provided the default OpenAI endpoint will be used.
*/
url?: string;
/**
* The key for the model. If 'true' is provided the global OpenAI API key will be used.
*/
apiKey: string | true | undefined;
}
export interface OpenAiLanguageModelsManager {
apiKey: string | undefined;
Expand Down
8 changes: 4 additions & 4 deletions packages/ai-openai/src/node/openai-language-model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ export class OpenAiModel implements LanguageModel {
* @param model the model id as it is used by the OpenAI API
* @param openAIInitializer initializer for the OpenAI client, used for each request.
*/
constructor(public readonly id: string, public model: string, protected apiKey: (() => string | undefined) | undefined, public url: string | undefined) { }
constructor(public readonly id: string, public model: string, public apiKey: () => string | undefined, public url: string | undefined) { }

async request(request: LanguageModelRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse> {
const openai = this.initializeOpenAi();
Expand Down Expand Up @@ -180,11 +180,11 @@ export class OpenAiModel implements LanguageModel {
}

protected initializeOpenAi(): OpenAI {
const apiKey = this.apiKey && this.apiKey();
const apiKey = this.apiKey();
if (!apiKey && !(this.url)) {
throw new Error('Please provide OPENAI_API_KEY in preferences or via environment variable');
}
// do not hand over API key to custom urls
return new OpenAI({ apiKey: this.url ? 'no-key' : apiKey, baseURL: this.url });
// We need to hand over "some" key, even if a custom url is not key protected as otherwise the OpenAI client will throw an error
return new OpenAI({ apiKey: apiKey ?? 'no-key', baseURL: this.url });
}
}
21 changes: 13 additions & 8 deletions packages/ai-openai/src/node/openai-language-models-manager-impl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,15 @@ export class OpenAiLanguageModelsManagerImpl implements OpenAiLanguageModelsMana
async createOrUpdateLanguageModels(...modelDescriptions: OpenAiModelDescription[]): Promise<void> {
for (const modelDescription of modelDescriptions) {
const model = await this.languageModelRegistry.getLanguageModel(modelDescription.id);
const apiKeyProvider = () => {
if (modelDescription.apiKey === true) {
return this.apiKey;
}
if (modelDescription.apiKey) {
return modelDescription.apiKey;
}
return undefined;
};
if (model) {
if (!(model instanceof OpenAiModel)) {
console.warn(`Open AI: model ${modelDescription.id} is not an OpenAI model`);
Expand All @@ -46,15 +55,11 @@ export class OpenAiLanguageModelsManagerImpl implements OpenAiLanguageModelsMana
console.info(`Open AI: skip creating model ${modelDescription.id} because it already exists`);
continue;
}
if (model.url !== modelDescription.url || model.model !== modelDescription.model) {
model.url = modelDescription.url;
model.model = modelDescription.model;
} else {
// This can happen during the initializing of more than one frontends.
console.info(`Open AI: skip creating or updating model ${modelDescription.id} because it already exists and is up to date`);
}
model.url = modelDescription.url;
model.model = modelDescription.model;
model.apiKey = apiKeyProvider;
} else {
this.languageModelRegistry.addLanguageModels([new OpenAiModel(modelDescription.id, modelDescription.model, () => this.apiKey, modelDescription.url)]);
this.languageModelRegistry.addLanguageModels([new OpenAiModel(modelDescription.id, modelDescription.model, apiKeyProvider, modelDescription.url)]);
}
}
}
Expand Down

0 comments on commit fc1b88e

Please sign in to comment.