|
| 1 | +/*--------------------------------------------------------------------------------------------- |
| 2 | + * Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | + * Licensed under the MIT License. See License.txt in the project root for license information. |
| 4 | + *--------------------------------------------------------------------------------------------*/ |
| 5 | + |
| 6 | +import assert from 'assert'; |
| 7 | +import { CancellationToken } from '../../../../../../base/common/cancellation.js'; |
| 8 | +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; |
| 9 | +import { Position } from '../../../../../../editor/common/core/position.js'; |
| 10 | +import { CompletionContext, CompletionTriggerKind } from '../../../../../../editor/common/languages.js'; |
| 11 | +import { ContextKeyService } from '../../../../../../platform/contextkey/browser/contextKeyService.js'; |
| 12 | +import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js'; |
| 13 | +import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; |
| 14 | +import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; |
| 15 | +import { workbenchInstantiationService } from '../../../../../test/browser/workbenchTestServices.js'; |
| 16 | +import { LanguageModelToolsService } from '../../../browser/languageModelToolsService.js'; |
| 17 | +import { ChatConfiguration } from '../../../common/constants.js'; |
| 18 | +import { ILanguageModelToolsService, IToolData, ToolDataSource } from '../../../common/languageModelToolsService.js'; |
| 19 | +import { PromptBodyAutocompletion } from '../../../common/promptSyntax/languageProviders/promptBodyAutocompletion.js'; |
| 20 | +import { createTextModel } from '../../../../../../editor/test/common/testTextModel.js'; |
| 21 | +import { URI } from '../../../../../../base/common/uri.js'; |
| 22 | +import { getLanguageIdForPromptsType, PromptsType } from '../../../common/promptSyntax/promptTypes.js'; |
| 23 | +import { getPromptFileExtension } from '../../../common/promptSyntax/config/promptFileLocations.js'; |
| 24 | +import { IFileService } from '../../../../../../platform/files/common/files.js'; |
| 25 | +import { FileService } from '../../../../../../platform/files/common/fileService.js'; |
| 26 | +import { VSBuffer } from '../../../../../../base/common/buffer.js'; |
| 27 | +import { InMemoryFileSystemProvider } from '../../../../../../platform/files/common/inMemoryFilesystemProvider.js'; |
| 28 | +import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js'; |
| 29 | +import { Range } from '../../../../../../editor/common/core/range.js'; |
| 30 | + |
| 31 | +suite('PromptBodyAutocompletion', () => { |
| 32 | + const disposables = ensureNoDisposablesAreLeakedInTestSuite(); |
| 33 | + |
| 34 | + let instaService: TestInstantiationService; |
| 35 | + let completionProvider: PromptBodyAutocompletion; |
| 36 | + |
| 37 | + setup(async () => { |
| 38 | + const testConfigService = new TestConfigurationService(); |
| 39 | + testConfigService.setUserConfiguration(ChatConfiguration.ExtensionToolsEnabled, true); |
| 40 | + instaService = workbenchInstantiationService({ |
| 41 | + contextKeyService: () => disposables.add(new ContextKeyService(testConfigService)), |
| 42 | + configurationService: () => testConfigService |
| 43 | + }, disposables); |
| 44 | + instaService.stub(ILogService, new NullLogService()); |
| 45 | + const fileService = disposables.add(instaService.createInstance(FileService)); |
| 46 | + instaService.stub(IFileService, fileService); |
| 47 | + |
| 48 | + const fileSystemProvider = disposables.add(new InMemoryFileSystemProvider()); |
| 49 | + disposables.add(fileService.registerProvider('test', fileSystemProvider)); |
| 50 | + |
| 51 | + // Create some test files and directories |
| 52 | + await fileService.createFolder(URI.parse('test:///workspace')); |
| 53 | + await fileService.createFolder(URI.parse('test:///workspace/src')); |
| 54 | + await fileService.createFolder(URI.parse('test:///workspace/docs')); |
| 55 | + await fileService.writeFile(URI.parse('test:///workspace/src/index.ts'), VSBuffer.fromString('export function hello() {}')); |
| 56 | + await fileService.writeFile(URI.parse('test:///workspace/README.md'), VSBuffer.fromString('# Project')); |
| 57 | + await fileService.writeFile(URI.parse('test:///workspace/package.json'), VSBuffer.fromString('{}')); |
| 58 | + |
| 59 | + const toolService = disposables.add(instaService.createInstance(LanguageModelToolsService)); |
| 60 | + |
| 61 | + const testTool1 = { id: 'testTool1', displayName: 'tool1', canBeReferencedInPrompt: true, modelDescription: 'Test Tool 1', source: ToolDataSource.External, inputSchema: {} } satisfies IToolData; |
| 62 | + disposables.add(toolService.registerToolData(testTool1)); |
| 63 | + |
| 64 | + const testTool2 = { id: 'testTool2', displayName: 'tool2', canBeReferencedInPrompt: true, toolReferenceName: 'tool2', modelDescription: 'Test Tool 2', source: ToolDataSource.External, inputSchema: {} } satisfies IToolData; |
| 65 | + disposables.add(toolService.registerToolData(testTool2)); |
| 66 | + |
| 67 | + const myExtSource = { type: 'extension', label: 'My Extension', extensionId: new ExtensionIdentifier('My.extension') } satisfies ToolDataSource; |
| 68 | + const testTool3 = { id: 'testTool3', displayName: 'tool3', canBeReferencedInPrompt: true, toolReferenceName: 'tool3', modelDescription: 'Test Tool 3', source: myExtSource, inputSchema: {} } satisfies IToolData; |
| 69 | + disposables.add(toolService.registerToolData(testTool3)); |
| 70 | + |
| 71 | + const prExtSource = { type: 'extension', label: 'GitHub Pull Request Extension', extensionId: new ExtensionIdentifier('github.vscode-pull-request-github') } satisfies ToolDataSource; |
| 72 | + const prExtTool1 = { id: 'suggestFix', canBeReferencedInPrompt: true, toolReferenceName: 'suggest-fix', modelDescription: 'tool4', displayName: 'Test Tool 4', source: prExtSource, inputSchema: {} } satisfies IToolData; |
| 73 | + disposables.add(toolService.registerToolData(prExtTool1)); |
| 74 | + |
| 75 | + instaService.set(ILanguageModelToolsService, toolService); |
| 76 | + |
| 77 | + completionProvider = instaService.createInstance(PromptBodyAutocompletion); |
| 78 | + }); |
| 79 | + |
| 80 | + async function getCompletions(content: string, line: number, column: number, promptType: PromptsType) { |
| 81 | + const languageId = getLanguageIdForPromptsType(promptType); |
| 82 | + const model = disposables.add(createTextModel(content, languageId, undefined, URI.parse('test://workspace/test' + getPromptFileExtension(promptType)))); |
| 83 | + const position = new Position(line, column); |
| 84 | + const context: CompletionContext = { triggerKind: CompletionTriggerKind.Invoke }; |
| 85 | + const result = await completionProvider.provideCompletionItems(model, position, context, CancellationToken.None); |
| 86 | + if (!result || !result.suggestions) { |
| 87 | + return []; |
| 88 | + } |
| 89 | + const lineContent = model.getLineContent(position.lineNumber); |
| 90 | + return result.suggestions.map(s => { |
| 91 | + assert(s.range instanceof Range); |
| 92 | + return { |
| 93 | + label: s.label, |
| 94 | + result: lineContent.substring(0, s.range.startColumn - 1) + s.insertText + lineContent.substring(s.range.endColumn - 1) |
| 95 | + }; |
| 96 | + }); |
| 97 | + } |
| 98 | + |
| 99 | + suite('prompt body completions', () => { |
| 100 | + test('default suggestions', async () => { |
| 101 | + const content = [ |
| 102 | + '---', |
| 103 | + 'description: "Test"', |
| 104 | + '---', |
| 105 | + '', |
| 106 | + 'Use # to reference a file or tool.', |
| 107 | + 'One more #to' |
| 108 | + ].join('\n'); |
| 109 | + |
| 110 | + { |
| 111 | + const actual = (await getCompletions(content, 5, 6, PromptsType.prompt)); |
| 112 | + assert.deepEqual(actual, [ |
| 113 | + { |
| 114 | + label: 'file:', |
| 115 | + result: 'Use #file: to reference a file or tool.' |
| 116 | + }, |
| 117 | + { |
| 118 | + label: 'tool:', |
| 119 | + result: 'Use #tool: to reference a file or tool.' |
| 120 | + } |
| 121 | + ]); |
| 122 | + } |
| 123 | + { |
| 124 | + const actual = (await getCompletions(content, 6, 13, PromptsType.prompt)); |
| 125 | + assert.deepEqual(actual, [ |
| 126 | + { |
| 127 | + label: 'file:', |
| 128 | + result: 'One more #file:' |
| 129 | + }, |
| 130 | + { |
| 131 | + label: 'tool:', |
| 132 | + result: 'One more #tool:' |
| 133 | + } |
| 134 | + ]); |
| 135 | + } |
| 136 | + }); |
| 137 | + |
| 138 | + test('tool suggestions', async () => { |
| 139 | + const content = [ |
| 140 | + '---', |
| 141 | + 'description: "Test"', |
| 142 | + '---', |
| 143 | + '', |
| 144 | + 'Use #tool: to reference a tool.', |
| 145 | + ].join('\n'); |
| 146 | + { |
| 147 | + const actual = (await getCompletions(content, 5, 11, PromptsType.prompt)); |
| 148 | + assert.deepEqual(actual, [ |
| 149 | + { |
| 150 | + label: 'tool1', |
| 151 | + result: 'Use #tool:tool1 to reference a tool.' |
| 152 | + }, |
| 153 | + { |
| 154 | + label: 'tool2', |
| 155 | + result: 'Use #tool:tool2 to reference a tool.' |
| 156 | + }, |
| 157 | + { |
| 158 | + label: 'my.extension/tool3', |
| 159 | + result: 'Use #tool:my.extension/tool3 to reference a tool.' |
| 160 | + }, |
| 161 | + { |
| 162 | + label: 'github.vscode-pull-request-github/suggest-fix', |
| 163 | + result: 'Use #tool:github.vscode-pull-request-github/suggest-fix to reference a tool.' |
| 164 | + } |
| 165 | + ]); |
| 166 | + } |
| 167 | + }); |
| 168 | + }); |
| 169 | +}); |
0 commit comments