From 2d3068925187e85044b207a8e7c61e2c976a5919 Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Fri, 8 Nov 2024 09:46:56 +0100 Subject: [PATCH 1/5] Improves the relevances of codestral completion --- package.json | 1 + src/completion-provider.ts | 6 ++ src/index.ts | 5 +- src/llm-models/base-completer.ts | 5 ++ src/llm-models/codestral-completer.ts | 81 ++++++++++++++++++++++----- src/provider.ts | 7 ++- yarn.lock | 3 +- 7 files changed, 91 insertions(+), 17 deletions(-) diff --git a/package.json b/package.json index 539736a..d70c7ff 100644 --- a/package.json +++ b/package.json @@ -62,6 +62,7 @@ "@jupyterlab/settingregistry": "^4.2.0", "@langchain/core": "^0.3.13", "@langchain/mistralai": "^0.1.1", + "@lumino/commands": "^2.1.2", "@lumino/coreutils": "^2.1.2", "@lumino/polling": "^2.1.2", "@lumino/signaling": "^2.1.2" diff --git a/src/completion-provider.ts b/src/completion-provider.ts index e7000c5..f036c0e 100644 --- a/src/completion-provider.ts +++ b/src/completion-provider.ts @@ -16,6 +16,7 @@ export class CompletionProvider implements IInlineCompletionProvider { constructor(options: CompletionProvider.IOptions) { const { name, settings } = options; + this._requestCompletion = options.requestCompletion; this.setCompleter(name, settings); } @@ -28,6 +29,9 @@ export class CompletionProvider implements IInlineCompletionProvider { setCompleter(name: string, settings: ReadonlyPartialJSONObject) { try { this._completer = getCompleter(name, settings); + if (this._completer) { + this._completer.requestCompletion = this._requestCompletion; + } this._name = this._completer === null ? 'None' : name; } catch (e: any) { this._completer = null; @@ -65,11 +69,13 @@ export class CompletionProvider implements IInlineCompletionProvider { } private _name: string = 'None'; + private _requestCompletion: () => void; private _completer: IBaseCompleter | null = null; } export namespace CompletionProvider { export interface IOptions extends BaseCompleter.IOptions { name: string; + requestCompletion: () => void; } } diff --git a/src/index.ts b/src/index.ts index 76d2ab2..b056487 100644 --- a/src/index.ts +++ b/src/index.ts @@ -100,7 +100,10 @@ const aiProviderPlugin: JupyterFrontEndPlugin = { manager: ICompletionProviderManager, settingRegistry: ISettingRegistry ): IAIProvider => { - const aiProvider = new AIProvider({ completionProviderManager: manager }); + const aiProvider = new AIProvider({ + completionProviderManager: manager, + requestCompletion: () => app.commands.execute('inline-completer:invoke') + }); settingRegistry .load(aiProviderPlugin.id) diff --git a/src/llm-models/base-completer.ts b/src/llm-models/base-completer.ts index fb84f4f..0828a9c 100644 --- a/src/llm-models/base-completer.ts +++ b/src/llm-models/base-completer.ts @@ -11,6 +11,11 @@ export interface IBaseCompleter { */ provider: LLM; + /** + * The function to fetch a new completion. + */ + requestCompletion?: () => void; + /** * The fetch request for the LLM completer. */ diff --git a/src/llm-models/codestral-completer.ts b/src/llm-models/codestral-completer.ts index efa7934..c21097a 100644 --- a/src/llm-models/codestral-completer.ts +++ b/src/llm-models/codestral-completer.ts @@ -16,27 +16,70 @@ const INTERVAL = 1000; export class CodestralCompleter implements IBaseCompleter { constructor(options: BaseCompleter.IOptions) { + // this._requestCompletion = options.requestCompletion; this._mistralProvider = new MistralAI({ ...options.settings }); - this._throttler = new Throttler(async (data: CompletionRequest) => { - const response = await this._mistralProvider.completionWithRetry( - data, - {}, - false - ); - const items = response.choices.map((choice: any) => { - return { insertText: choice.message.content as string }; - }); + this._throttler = new Throttler( + async (data: CompletionRequest) => { + this._invokedData = data; + let fetchAgain = false; - return { - items - }; - }, INTERVAL); + // Request completion. + const response = await this._mistralProvider.completionWithRetry( + data, + {}, + false + ); + + // Extract results of completion request. + let items = response.choices.map((choice: any) => { + return { insertText: choice.message.content as string }; + }); + + // Check if the prompt has changed during the request. + if (this._invokedData.prompt !== this._currentData?.prompt) { + // The current prompt does not include the invoked one, the result is + // cancelled and a new completion will be requested. + if (!this._currentData?.prompt.startsWith(this._invokedData.prompt)) { + fetchAgain = true; + items = []; + } else { + // Check if some results contain the current prompt, and return them if so, + // otherwise request completion again. + const newItems: { insertText: string }[] = []; + items.forEach(item => { + const result = this._invokedData!.prompt + item.insertText; + if (result.startsWith(this._currentData!.prompt)) { + const insertText = result.slice( + this._currentData!.prompt.length + ); + newItems.push({ insertText }); + } + }); + if (newItems.length) { + items = newItems; + } else { + fetchAgain = true; + items = []; + } + } + } + return { + items, + fetchAgain + }; + }, + { limit: INTERVAL } + ); } get provider(): LLM { return this._mistralProvider; } + set requestCompletion(value: () => void) { + this._requestCompletion = value; + } + async fetch( request: CompletionHandler.IRequest, context: IInlineCompletionContext @@ -59,13 +102,23 @@ export class CodestralCompleter implements IBaseCompleter { }; try { - return this._throttler.invoke(data); + this._currentData = data; + const completionResult = await this._throttler.invoke(data); + if (completionResult.fetchAgain) { + if (this._requestCompletion) { + this._requestCompletion(); + } + } + return { items: completionResult.items }; } catch (error) { console.error('Error fetching completions', error); return { items: [] }; } } + private _requestCompletion?: () => void; private _throttler: Throttler; private _mistralProvider: MistralAI; + private _invokedData: CompletionRequest | null = null; + private _currentData: CompletionRequest | null = null; } diff --git a/src/provider.ts b/src/provider.ts index 6019785..1347b5b 100644 --- a/src/provider.ts +++ b/src/provider.ts @@ -12,7 +12,8 @@ export class AIProvider implements IAIProvider { constructor(options: AIProvider.IOptions) { this._completionProvider = new CompletionProvider({ name: 'None', - settings: {} + settings: {}, + requestCompletion: options.requestCompletion }); options.completionProviderManager.registerInlineProvider( this._completionProvider @@ -103,6 +104,10 @@ export namespace AIProvider { * The completion provider manager in which register the LLM completer. */ completionProviderManager: ICompletionProviderManager; + /** + * The application commands registry. + */ + requestCompletion: () => void; } /** diff --git a/yarn.lock b/yarn.lock index 47e6599..2278c09 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1740,7 +1740,7 @@ __metadata: languageName: node linkType: hard -"@lumino/commands@npm:^2.3.0": +"@lumino/commands@npm:^2.1.2, @lumino/commands@npm:^2.3.0": version: 2.3.0 resolution: "@lumino/commands@npm:2.3.0" dependencies: @@ -4885,6 +4885,7 @@ __metadata: "@jupyterlab/settingregistry": ^4.2.0 "@langchain/core": ^0.3.13 "@langchain/mistralai": ^0.1.1 + "@lumino/commands": ^2.1.2 "@lumino/coreutils": ^2.1.2 "@lumino/polling": ^2.1.2 "@lumino/signaling": ^2.1.2 From 23932d2aabec2d6a0a887b109a1f2b0bcbb1855f Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Fri, 8 Nov 2024 10:36:18 +0100 Subject: [PATCH 2/5] Add a timeout to avoid endless requests --- src/llm-models/codestral-completer.ts | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/src/llm-models/codestral-completer.ts b/src/llm-models/codestral-completer.ts index c21097a..33e2e17 100644 --- a/src/llm-models/codestral-completer.ts +++ b/src/llm-models/codestral-completer.ts @@ -9,11 +9,16 @@ import { CompletionRequest } from '@mistralai/mistralai'; import { BaseCompleter, IBaseCompleter } from './base-completer'; -/* +/** * The Mistral API has a rate limit of 1 request per second */ const INTERVAL = 1000; +/** + * Timeout to avoid endless requests + */ +const REQUEST_TIMEOUT = 3000; + export class CodestralCompleter implements IBaseCompleter { constructor(options: BaseCompleter.IOptions) { // this._requestCompletion = options.requestCompletion; @@ -24,11 +29,22 @@ export class CodestralCompleter implements IBaseCompleter { let fetchAgain = false; // Request completion. - const response = await this._mistralProvider.completionWithRetry( + const request = this._mistralProvider.completionWithRetry( data, {}, false ); + const timeoutPromise = new Promise(resolve => { + return setTimeout(() => resolve(null), REQUEST_TIMEOUT); + }); + + const response = await Promise.race([request, timeoutPromise]); + if (response === null) { + return { + items: [], + fetchAgain: true + }; + } // Extract results of completion request. let items = response.choices.map((choice: any) => { From 1304f4062d666377815250418185b95fb083300c Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Tue, 12 Nov 2024 09:51:36 +0100 Subject: [PATCH 3/5] Remove unused dependency --- package.json | 1 - yarn.lock | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/package.json b/package.json index d70c7ff..539736a 100644 --- a/package.json +++ b/package.json @@ -62,7 +62,6 @@ "@jupyterlab/settingregistry": "^4.2.0", "@langchain/core": "^0.3.13", "@langchain/mistralai": "^0.1.1", - "@lumino/commands": "^2.1.2", "@lumino/coreutils": "^2.1.2", "@lumino/polling": "^2.1.2", "@lumino/signaling": "^2.1.2" diff --git a/yarn.lock b/yarn.lock index 2278c09..47e6599 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1740,7 +1740,7 @@ __metadata: languageName: node linkType: hard -"@lumino/commands@npm:^2.1.2, @lumino/commands@npm:^2.3.0": +"@lumino/commands@npm:^2.3.0": version: 2.3.0 resolution: "@lumino/commands@npm:2.3.0" dependencies: @@ -4885,7 +4885,6 @@ __metadata: "@jupyterlab/settingregistry": ^4.2.0 "@langchain/core": ^0.3.13 "@langchain/mistralai": ^0.1.1 - "@lumino/commands": ^2.1.2 "@lumino/coreutils": ^2.1.2 "@lumino/polling": ^2.1.2 "@lumino/signaling": ^2.1.2 From 3e7142ac2f3b78bccbaceb388c24e5dddfea53db Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Mon, 2 Dec 2024 14:30:16 +0100 Subject: [PATCH 4/5] Fetch again if the prompt has changed between the request and the response --- src/llm-models/codestral-completer.ts | 38 +++------------------------ 1 file changed, 4 insertions(+), 34 deletions(-) diff --git a/src/llm-models/codestral-completer.ts b/src/llm-models/codestral-completer.ts index 33e2e17..77d5e86 100644 --- a/src/llm-models/codestral-completer.ts +++ b/src/llm-models/codestral-completer.ts @@ -25,8 +25,7 @@ export class CodestralCompleter implements IBaseCompleter { this._mistralProvider = new MistralAI({ ...options.settings }); this._throttler = new Throttler( async (data: CompletionRequest) => { - this._invokedData = data; - let fetchAgain = false; + const invokedData = data; // Request completion. const request = this._mistralProvider.completionWithRetry( @@ -38,8 +37,9 @@ export class CodestralCompleter implements IBaseCompleter { 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) { + if (response === null || invokedData.prompt !== this._currentData?.prompt) { return { items: [], fetchAgain: true @@ -51,37 +51,8 @@ export class CodestralCompleter implements IBaseCompleter { return { insertText: choice.message.content as string }; }); - // Check if the prompt has changed during the request. - if (this._invokedData.prompt !== this._currentData?.prompt) { - // The current prompt does not include the invoked one, the result is - // cancelled and a new completion will be requested. - if (!this._currentData?.prompt.startsWith(this._invokedData.prompt)) { - fetchAgain = true; - items = []; - } else { - // Check if some results contain the current prompt, and return them if so, - // otherwise request completion again. - const newItems: { insertText: string }[] = []; - items.forEach(item => { - const result = this._invokedData!.prompt + item.insertText; - if (result.startsWith(this._currentData!.prompt)) { - const insertText = result.slice( - this._currentData!.prompt.length - ); - newItems.push({ insertText }); - } - }); - if (newItems.length) { - items = newItems; - } else { - fetchAgain = true; - items = []; - } - } - } return { - items, - fetchAgain + items }; }, { limit: INTERVAL } @@ -135,6 +106,5 @@ export class CodestralCompleter implements IBaseCompleter { private _requestCompletion?: () => void; private _throttler: Throttler; private _mistralProvider: MistralAI; - private _invokedData: CompletionRequest | null = null; private _currentData: CompletionRequest | null = null; } From cc2b170998f1eecbf1b60ec81f6120e32d7ebcd7 Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Mon, 2 Dec 2024 14:56:57 +0100 Subject: [PATCH 5/5] lint --- src/llm-models/codestral-completer.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/llm-models/codestral-completer.ts b/src/llm-models/codestral-completer.ts index 77d5e86..4db7313 100644 --- a/src/llm-models/codestral-completer.ts +++ b/src/llm-models/codestral-completer.ts @@ -39,7 +39,10 @@ export class CodestralCompleter implements IBaseCompleter { // 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 ( + response === null || + invokedData.prompt !== this._currentData?.prompt + ) { return { items: [], fetchAgain: true @@ -47,7 +50,7 @@ export class CodestralCompleter implements IBaseCompleter { } // Extract results of completion request. - let items = response.choices.map((choice: any) => { + const items = response.choices.map((choice: any) => { return { insertText: choice.message.content as string }; });