From 415dd0470c33a1eb49283c71175f4e69c80cfbe5 Mon Sep 17 00:00:00 2001 From: Oleg Solomko Date: Sun, 22 Dec 2024 14:03:10 -0800 Subject: [PATCH] wip --- .../contrib/chat/common/basePromptParser.ts | 6 +- .../contrib/chat/common/basePromptTypes.d.ts | 9 +- .../chat/common/filePromptContentProvider.ts | 18 +- .../chat/common/promptContentsProviderBase.ts | 4 +- .../chat/common/textModelContentsProvider.ts | 4 +- .../chat/common/textModelPromptParser.ts | 2 +- .../test/common/textModelPromptParser.test.ts | 236 ++++++++++++++++-- 7 files changed, 236 insertions(+), 43 deletions(-) diff --git a/src/vs/workbench/contrib/chat/common/basePromptParser.ts b/src/vs/workbench/contrib/chat/common/basePromptParser.ts index 18c09ad574a..9d01d0793ec 100644 --- a/src/vs/workbench/contrib/chat/common/basePromptParser.ts +++ b/src/vs/workbench/contrib/chat/common/basePromptParser.ts @@ -5,7 +5,7 @@ // TODO: @legomushroom - cleanup import { URI } from '../../../../base/common/uri.js'; -import { IPromptFileReference, IPromptProvider, TPromptPart } from './basePromptTypes.js'; +import { IPromptFileReference, IPromptContentsProvider, TPromptPart } from './basePromptTypes.js'; import { assert, assertNever } from '../../../../base/common/assert.js'; import { Emitter } from '../../../../base/common/event.js'; import { BaseDecoder } from '../../../../base/common/codecs/baseDecoder.js'; @@ -79,7 +79,7 @@ export class PromptLine extends Disposable { if (token instanceof FileReference) { const fileReference = this.instantiationService - .createInstance(PromptFileReference, token, dirname, seenReferences); + .createInstance(PromptFileReference, token, dirname, [...seenReferences]); this._tokens.push(fileReference); @@ -164,7 +164,7 @@ const PROMPT_SNIPPETS_CONFIG_KEY: string = 'chat.experimental.prompt-snippets'; /** * TODO: @legomushroom */ -export class BasePromptParser extends Disposable { +export class BasePromptParser extends Disposable { public disposed: boolean = false; /** diff --git a/src/vs/workbench/contrib/chat/common/basePromptTypes.d.ts b/src/vs/workbench/contrib/chat/common/basePromptTypes.d.ts index 11c69213928..e66460e3ffc 100644 --- a/src/vs/workbench/contrib/chat/common/basePromptTypes.d.ts +++ b/src/vs/workbench/contrib/chat/common/basePromptTypes.d.ts @@ -19,14 +19,15 @@ type TOnContentChangedCallback = (streamOrError: ReadableStream | ParseErr /** * TODO: @legomushroom - move to the correct place */ -export interface IPromptProvider extends IDisposable { +export interface IPromptContentsProvider extends IDisposable { start(): void; onContentChanged(callback: TOnContentChangedCallback): IDisposable; readonly uri: URI; } +type TPromptPartTypes = 'file-reference'; export interface IPromptPart { - readonly type: 'file-reference'; + readonly type: TPromptPartTypes; readonly range: Range; readonly text: string; } @@ -35,11 +36,11 @@ export interface IPromptFileReference extends IPromptPart { readonly type: 'file-reference'; readonly uri: URI; readonly path: string; + readonly resolveFailed: boolean | undefined; + readonly errorCondition: ParseError | undefined; tokensTree: readonly TPromptPart[]; allValidFileReferenceUris: readonly URI[]; validFileReferences: readonly IPromptFileReference[]; - readonly resolveFailed: boolean | undefined; - readonly errorCondition: ParseError | undefined; } /** diff --git a/src/vs/workbench/contrib/chat/common/filePromptContentProvider.ts b/src/vs/workbench/contrib/chat/common/filePromptContentProvider.ts index 9abc8bc0963..640709b6884 100644 --- a/src/vs/workbench/contrib/chat/common/filePromptContentProvider.ts +++ b/src/vs/workbench/contrib/chat/common/filePromptContentProvider.ts @@ -4,9 +4,9 @@ *--------------------------------------------------------------------------------------------*/ import { URI } from '../../../../base/common/uri.js'; -import { IPromptProvider } from './basePromptTypes.js'; +import { IPromptContentsProvider } from './basePromptTypes.js'; import { BasePromptParser } from './basePromptParser.js'; -// import { cancelOnDispose } from './cancelOnDisposeDecorator.js'; +import { cancelOnDispose } from './cancelOnDisposeDecorator.js'; import { basename } from '../../../../base/common/resources.js'; import { CancellationError } from '../../../../base/common/errors.js'; import { CancellationToken } from '../../../../base/common/cancellation.js'; @@ -21,7 +21,7 @@ import { FileChangesEvent, FileChangeType, IFileService } from '../../../../plat * TODO: @legomushroom * TODO: @legomushroom - move to the correct place */ -export class FilePromptContentProvider extends PromptContentsProviderBase implements IPromptProvider { +export class FilePromptContentProvider extends PromptContentsProviderBase implements IPromptContentsProvider { /** * TODO: @legomushroom */ @@ -68,7 +68,7 @@ export class FilePromptContentProvider extends PromptContentsProviderBase, -> extends Disposable implements IPromptProvider { +> extends Disposable implements IPromptContentsProvider { /** * Internal event emitter for the prompt contents change event. Classes that extend * this abstract class are responsible to use this emitter to fire the contents change diff --git a/src/vs/workbench/contrib/chat/common/textModelContentsProvider.ts b/src/vs/workbench/contrib/chat/common/textModelContentsProvider.ts index 70ba3d26ee2..7f1271816bc 100644 --- a/src/vs/workbench/contrib/chat/common/textModelContentsProvider.ts +++ b/src/vs/workbench/contrib/chat/common/textModelContentsProvider.ts @@ -3,7 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { IPromptProvider } from './basePromptTypes.js'; +import { IPromptContentsProvider } from './basePromptTypes.js'; import { ITextModel } from '../../../../editor/common/model.js'; import { PromptContentsProviderBase } from './promptContentsProviderBase.js'; import { Line } from '../../../../editor/common/codecs/linesCodec/tokens/line.js'; @@ -14,7 +14,7 @@ import { IModelContentChangedEvent } from '../../../../editor/common/textModelEv * TODO: @legomushroom * TODO: @legomushroom - move to a correct place */ -export class TextModelContentsProvider extends PromptContentsProviderBase implements IPromptProvider { +export class TextModelContentsProvider extends PromptContentsProviderBase implements IPromptContentsProvider { public readonly uri = this.model.uri; constructor( diff --git a/src/vs/workbench/contrib/chat/common/textModelPromptParser.ts b/src/vs/workbench/contrib/chat/common/textModelPromptParser.ts index 72c33b8c0fa..d5e94b60f53 100644 --- a/src/vs/workbench/contrib/chat/common/textModelPromptParser.ts +++ b/src/vs/workbench/contrib/chat/common/textModelPromptParser.ts @@ -27,6 +27,6 @@ export class TextModelPromptParser extends BasePromptParser => { return new Promise(resolve => setTimeout(resolve, ms)); }; +/** + * Helper function that allows to await for a random amount of time. + * @param maxMs The `maximum` amount of time to wait, in milliseconds. + * @param minMs [`optional`] The `minimum` amount of time to wait, in milliseconds. + * TODO: @legomushroom - move to a common utility + */ +const waitRandom = (maxMs: number, minMs: number = 0): Promise => { + return wait(randomInt(maxMs, minMs)); +}; + +/** + * Represents a file reference with an expected + * error condition value for testing purposes. + * TODO: @legomushroom - move to a common utility + */ +class ExpectedReference { + /** + * URI component of the expected reference. + */ + public readonly uri: URI; + + constructor( + dirname: URI, + public readonly lineToken: FileReference, + public readonly errorCondition?: TErrorCondition, + ) { + this.uri = extUri.resolvePath(dirname, lineToken.path); + } + + /** + * String representation of the expected reference. + */ + public toString(): string { + return `#file:${this.uri.path}`; + } +} + + +/** + * A reusable test utility to test the `TextModelParser` class. + */ +class TestTextModelParser extends Disposable { + /** + * Reference to the test instance of the the `TextModel`. + */ + private readonly model: ITextModel; + + constructor( + modelUri: URI, + private readonly modelContent: string, + private readonly expectedReferences: ExpectedReference[], + @IInstantiationService private readonly initService: IInstantiationService, + @IFileService private readonly fileService: IFileService, + ) { + super(); + + // create in-memory file system + const fileSystemProvider = this._register(new InMemoryFileSystemProvider()); + this._register(this.fileService.registerProvider(Schemas.file, fileSystemProvider)); + + this.model = this._register(instantiateTextModel( + initService, + this.modelContent, + null, + undefined, + modelUri, + )); + } + + /** + * Directory containing the test model. + */ + public get dirname(): URI { + return extUri.dirname(this.model.uri); + } + + /** + * Run the test. + */ + public async run() { + + // TODO: @legomushroom - add description + // TODO: @legomushroom - test without the delay + await waitRandom(5); + + // start resolving references for the specified root file + const parser = this._register( + this.initService.createInstance( + TextModelPromptParser, + this.model, + [], + ), + ).start(); + + // TODO: @legomushroom - add description + await wait(100); + + // resolve the root file reference including all nested references + const resolvedReferences: readonly (IPromptFileReference | undefined)[] = parser.tokensTree; + + for (let i = 0; i < this.expectedReferences.length; i++) { + const expectedReference = this.expectedReferences[i]; + const resolvedReference = resolvedReferences[i]; + + assert( + (resolvedReference) && + (resolvedReference.uri.toString() === expectedReference.uri.toString()), + [ + `Expected ${i}th resolved reference URI to be '${expectedReference.uri}'`, + `got '${resolvedReference?.uri}'.`, + ].join(', '), + ); + + // assert( + // resolvedReference.token.equals(expectedReference.lineToken), + // [ + // `Expected ${i}th resolved reference token to be '${expectedReference.lineToken}'`, + // `got '${resolvedReference.token}'.`, + // ].join(', '), + // ); + + if (expectedReference.errorCondition === undefined) { + assert( + resolvedReference.errorCondition === undefined, + [ + `Expected ${i}th error condition to be 'undefined'`, + `got '${resolvedReference.errorCondition}'.`, + ].join(', '), + ); + continue; + } + + assert( + expectedReference.errorCondition.equal(resolvedReference.errorCondition), + [ + `Expected ${i}th error condition to be '${expectedReference.errorCondition}'`, + `got '${resolvedReference.errorCondition}'.`, + ].join(', '), + ); + } + + assert.strictEqual( + resolvedReferences.length, + this.expectedReferences.length, + [ + `\nExpected(${this.expectedReferences.length}): [\n ${this.expectedReferences.join('\n ')}\n]`, + `Received(${resolvedReferences.length}): [\n ${resolvedReferences.join('\n ')}\n]`, + ].join('\n') + ); + } +} + +/** + * Create expected file reference for testing purposes. + * + * @param filePath The expected path of the file reference (without the `#file:` prefix). + * @param lineNumber The expected line number of the file reference. + * @param startColumnNumber The expected start column number of the file reference. + * TODO: @legomushroom - move to a common utility + */ +const createTestFileReference = ( + filePath: string, + lineNumber: number, + startColumnNumber: number, +): FileReference => { + const range = new Range( + lineNumber, + startColumnNumber, + lineNumber, + startColumnNumber + `#file:${filePath}`.length, + ); + + return new FileReference(range, filePath); +}; + suite('TextModelPromptParser (Unix)', function () { const testDisposables = ensureNoDisposablesAreLeakedInTestSuite(); @@ -30,34 +217,39 @@ suite('TextModelPromptParser (Unix)', function () { this.skip(); } - const disposables = new DisposableStore(); + const disposables = testDisposables.add(new DisposableStore()); const instantiationService = createModelServices(disposables); - const model = instantiateTextModel(instantiationService, 'My First Line\r\nMy \t#file:./some-file.md Line\r\nMy Third Line'); - model.registerDisposable(disposables); const nullLogService = testDisposables.add(new NullLogService()); const nullFileService = testDisposables.add(new FileService(nullLogService)); - instantiationService.stub(IFileService, nullFileService); - testDisposables.add(model); + const rootUri = URI.file('/root'); - await wait(100); + const test = testDisposables.add(instantiationService.createInstance( + TestTextModelParser, + extUri.joinPath(rootUri, 'file.txt'), + `My First Line\r\nMy \t#file:./some-file.md Line\t\t some content\t#file:/root/some-file.md\r\nMy Third Line`, + [ + new ExpectedReference( + rootUri, + createTestFileReference('./some-file.md', 2, 4), + new FileOpenFailed( + URI.joinPath(rootUri, './some-file.md'), + 'Failed to open file', + ), + ), + new ExpectedReference( + rootUri, + createTestFileReference('./some-file.md', 2, 4), + new FileOpenFailed( + URI.joinPath(rootUri, './some-file.md'), + 'Failed to open file', + ), + ), + ], + )); - const parser = testDisposables.add( - instantiationService.createInstance(TextModelPromptParser, model, []), - ); - - parser.start(); - - // TODO: @legomushroom - add description - await wait(100); - - const tokens = parser.tokens; - - assert( - tokens.length === 1, - `Expected 1 token, but got ${tokens.length}.`, - ); + await test.run(); }); });