From 0e65a7d704dec3429353e0dd58a80af96f297003 Mon Sep 17 00:00:00 2001 From: pmkrawczyk <63370626+pmkrawczyk@users.noreply.github.com> Date: Fri, 3 Jan 2025 18:55:39 +0100 Subject: [PATCH] Add external API support for Mistral completions Add optional external API endpoint for Mistral completions in jupyterlite extension. * Add `useExternalAPI` and `externalAPIEndpoint` settings in `schema/ai-provider.json`. * Modify `src/llm-models/codestral-completer.ts` to use the external API endpoint if specified in the settings. * Update `src/index.ts` to load the new settings from the settings registry. * Add `jupyter-lite.json` file to configure the new settings. --- jupyter-lite.json | 4 ++ schema/ai-provider.json | 12 ++++ src/index.ts | 4 +- src/llm-models/codestral-completer.ts | 92 +++++++++++++++++---------- 4 files changed, 78 insertions(+), 34 deletions(-) create mode 100644 jupyter-lite.json diff --git a/jupyter-lite.json b/jupyter-lite.json new file mode 100644 index 00000000..3d5d54c9 --- /dev/null +++ b/jupyter-lite.json @@ -0,0 +1,4 @@ +{ + "useExternalAPI": false, + "externalAPIEndpoint": "" +} diff --git a/schema/ai-provider.json b/schema/ai-provider.json index d4b9a049..0d9f5eee 100644 --- a/schema/ai-provider.json +++ b/schema/ai-provider.json @@ -15,6 +15,18 @@ "title": "The Codestral API key", "description": "The API key to use for Codestral", "default": "" + }, + "useExternalAPI": { + "type": "boolean", + "title": "Use External API", + "description": "Whether to use an external API endpoint for Mistral completions", + "default": false + }, + "externalAPIEndpoint": { + "type": "string", + "title": "External API Endpoint", + "description": "The external API endpoint URL for Mistral completions", + "default": "" } }, "additionalProperties": false diff --git a/src/index.ts b/src/index.ts index 56e39ed4..6ab38828 100644 --- a/src/index.ts +++ b/src/index.ts @@ -110,7 +110,9 @@ const aiProviderPlugin: JupyterFrontEndPlugin = { .then(settings => { const updateProvider = () => { const provider = settings.get('provider').composite as string; - aiProvider.setModels(provider, settings.composite); + const useExternalAPI = settings.get('useExternalAPI').composite as boolean; + const externalAPIEndpoint = settings.get('externalAPIEndpoint').composite as string; + aiProvider.setModels(provider, { ...settings.composite, useExternalAPI, externalAPIEndpoint }); }; settings.changed.connect(() => updateProvider()); diff --git a/src/llm-models/codestral-completer.ts b/src/llm-models/codestral-completer.ts index 4db73137..874df2ee 100644 --- a/src/llm-models/codestral-completer.ts +++ b/src/llm-models/codestral-completer.ts @@ -21,42 +21,71 @@ const REQUEST_TIMEOUT = 3000; export class CodestralCompleter implements IBaseCompleter { constructor(options: BaseCompleter.IOptions) { - // this._requestCompletion = options.requestCompletion; + this._useExternalAPI = options.settings.useExternalAPI as boolean; + this._externalAPIEndpoint = options.settings.externalAPIEndpoint as string; this._mistralProvider = new MistralAI({ ...options.settings }); this._throttler = new Throttler( async (data: CompletionRequest) => { const invokedData = data; - // Request completion. - const request = this._mistralProvider.completionWithRetry( - data, - {}, - false - ); - const timeoutPromise = new Promise(resolve => { - return setTimeout(() => resolve(null), REQUEST_TIMEOUT); - }); - - // Fetch again if the request is too long or if the prompt has changed. - const response = await Promise.race([request, timeoutPromise]); - if ( - response === null || - invokedData.prompt !== this._currentData?.prompt - ) { + if (this._useExternalAPI && this._externalAPIEndpoint) { + try { + const response = await fetch(this._externalAPIEndpoint, { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify(data) + }); + + if (!response.ok) { + throw new Error('Failed to fetch from external API'); + } + + const result = await response.json(); + const items = result.choices.map((choice: any) => { + return { insertText: choice.message.content as string }; + }); + + return { + items + }; + } catch (error) { + console.error('Error fetching completions from external API', error); + return { items: [] }; + } + } else { + // Request completion. + const request = this._mistralProvider.completionWithRetry( + data, + {}, + false + ); + const timeoutPromise = new Promise(resolve => { + return setTimeout(() => resolve(null), REQUEST_TIMEOUT); + }); + + // Fetch again if the request is too long or if the prompt has changed. + const response = await Promise.race([request, timeoutPromise]); + if ( + response === null || + invokedData.prompt !== this._currentData?.prompt + ) { + return { + items: [], + fetchAgain: true + }; + } + + // Extract results of completion request. + const items = response.choices.map((choice: any) => { + return { insertText: choice.message.content as string }; + }); + return { - items: [], - fetchAgain: true + items }; } - - // Extract results of completion request. - const items = response.choices.map((choice: any) => { - return { insertText: choice.message.content as string }; - }); - - return { - items - }; }, { limit: INTERVAL } ); @@ -82,12 +111,7 @@ export class CodestralCompleter implements IBaseCompleter { prompt, suffix, model: this._mistralProvider.model, - // temperature: 0, - // top_p: 1, - // max_tokens: 1024, - // min_tokens: 0, stream: false, - // random_seed: 1337, stop: [] }; @@ -110,4 +134,6 @@ export class CodestralCompleter implements IBaseCompleter { private _throttler: Throttler; private _mistralProvider: MistralAI; private _currentData: CompletionRequest | null = null; + private _useExternalAPI: boolean; + private _externalAPIEndpoint: string; }