|
| 1 | +/*! |
| 2 | + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +import assert from 'assert' |
| 7 | +import { |
| 8 | + ChatTriggerType, |
| 9 | + filePathSizeLimit, |
| 10 | + TriggerPayload, |
| 11 | + triggerPayloadToChatRequest, |
| 12 | +} from 'aws-core-vscode/codewhispererChat' |
| 13 | + |
| 14 | +describe('triggerPayloadToChatRequest', () => { |
| 15 | + const mockBasicPayload: TriggerPayload = { |
| 16 | + message: 'test message', |
| 17 | + filePath: 'test/path.ts', |
| 18 | + fileText: 'console.log("hello")', |
| 19 | + fileLanguage: 'typescript', |
| 20 | + additionalContents: [], |
| 21 | + relevantTextDocuments: [], |
| 22 | + useRelevantDocuments: false, |
| 23 | + customization: { arn: 'test:arn' }, |
| 24 | + trigger: ChatTriggerType.ChatMessage, |
| 25 | + contextLengths: { |
| 26 | + truncatedUserInputContextLength: 0, |
| 27 | + truncatedFocusFileContextLength: 0, |
| 28 | + truncatedWorkspaceContextLength: 0, |
| 29 | + truncatedAdditionalContextLengths: { |
| 30 | + promptContextLength: 0, |
| 31 | + ruleContextLength: 0, |
| 32 | + fileContextLength: 0, |
| 33 | + }, |
| 34 | + additionalContextLengths: { |
| 35 | + fileContextLength: 0, |
| 36 | + promptContextLength: 0, |
| 37 | + ruleContextLength: 0, |
| 38 | + }, |
| 39 | + workspaceContextLength: 0, |
| 40 | + userInputContextLength: 0, |
| 41 | + focusFileContextLength: 0, |
| 42 | + }, |
| 43 | + context: [], |
| 44 | + documentReferences: [], |
| 45 | + query: undefined, |
| 46 | + codeSelection: undefined, |
| 47 | + matchPolicy: undefined, |
| 48 | + codeQuery: undefined, |
| 49 | + userIntent: undefined, |
| 50 | + } |
| 51 | + |
| 52 | + const createLargeString = (size: number, prefix: string = '') => prefix + 'x'.repeat(size - prefix.length) |
| 53 | + |
| 54 | + const createPrompt = (size: number) => { |
| 55 | + return { |
| 56 | + name: 'prompt', |
| 57 | + description: 'prompt', |
| 58 | + relativePath: 'path-prompt', |
| 59 | + type: 'prompt', |
| 60 | + innerContext: createLargeString(size, 'prompt-'), |
| 61 | + } |
| 62 | + } |
| 63 | + |
| 64 | + const createRule = (size: number) => { |
| 65 | + return { |
| 66 | + name: 'rule', |
| 67 | + description: 'rule', |
| 68 | + relativePath: 'path-rule', |
| 69 | + type: 'rule', |
| 70 | + innerContext: createLargeString(size, 'rule-'), |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + const createFile = (size: number) => { |
| 75 | + return { |
| 76 | + name: 'file', |
| 77 | + description: 'file', |
| 78 | + relativePath: 'path-file', |
| 79 | + type: 'file', |
| 80 | + innerContext: createLargeString(size, 'file-'), |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + const createBaseTriggerPayload = (): TriggerPayload => ({ |
| 85 | + ...mockBasicPayload, |
| 86 | + message: '', |
| 87 | + fileText: '', |
| 88 | + filePath: 'test.ts', |
| 89 | + fileLanguage: 'typescript', |
| 90 | + customization: { arn: '' }, |
| 91 | + }) |
| 92 | + it('should convert basic trigger payload to chat request', () => { |
| 93 | + const result = triggerPayloadToChatRequest(mockBasicPayload) |
| 94 | + |
| 95 | + assert.notEqual(result, undefined) |
| 96 | + assert.strictEqual(result.conversationState.currentMessage?.userInputMessage?.content, 'test message') |
| 97 | + assert.strictEqual(result.conversationState.chatTriggerType, 'MANUAL') |
| 98 | + assert.strictEqual(result.conversationState.customizationArn, 'test:arn') |
| 99 | + }) |
| 100 | + |
| 101 | + it('should handle empty file path', () => { |
| 102 | + const emptyFilePathPayload = { |
| 103 | + ...mockBasicPayload, |
| 104 | + filePath: '', |
| 105 | + } |
| 106 | + |
| 107 | + const result = triggerPayloadToChatRequest(emptyFilePathPayload) |
| 108 | + |
| 109 | + assert.strictEqual( |
| 110 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document, |
| 111 | + undefined |
| 112 | + ) |
| 113 | + }) |
| 114 | + |
| 115 | + it('should filter out empty additional contents', () => { |
| 116 | + const payloadWithEmptyContents: TriggerPayload = { |
| 117 | + ...mockBasicPayload, |
| 118 | + additionalContents: [ |
| 119 | + { name: 'prompt1', description: 'prompt1', relativePath: 'path1', type: 'prompt', innerContext: '' }, |
| 120 | + { |
| 121 | + name: 'prompt2', |
| 122 | + description: 'prompt2', |
| 123 | + relativePath: 'path2', |
| 124 | + type: 'prompt', |
| 125 | + innerContext: 'valid content', |
| 126 | + }, |
| 127 | + ], |
| 128 | + } |
| 129 | + |
| 130 | + const result = triggerPayloadToChatRequest(payloadWithEmptyContents) |
| 131 | + |
| 132 | + assert.strictEqual( |
| 133 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext |
| 134 | + ?.length, |
| 135 | + 1 |
| 136 | + ) |
| 137 | + assert.strictEqual( |
| 138 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext.additionalContext?.[0] |
| 139 | + .innerContext, |
| 140 | + 'valid content' |
| 141 | + ) |
| 142 | + }) |
| 143 | + |
| 144 | + it('should handle unsupported programming language', () => { |
| 145 | + const unsupportedLanguagePayload = { |
| 146 | + ...mockBasicPayload, |
| 147 | + fileLanguage: 'unsupported', |
| 148 | + } |
| 149 | + |
| 150 | + const result = triggerPayloadToChatRequest(unsupportedLanguagePayload) |
| 151 | + |
| 152 | + assert.strictEqual( |
| 153 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document |
| 154 | + ?.programmingLanguage, |
| 155 | + undefined |
| 156 | + ) |
| 157 | + }) |
| 158 | + |
| 159 | + it('should truncate file path if it exceeds limit', () => { |
| 160 | + const longFilePath = 'a'.repeat(5000) |
| 161 | + const longFilePathPayload = { |
| 162 | + ...mockBasicPayload, |
| 163 | + filePath: longFilePath, |
| 164 | + } |
| 165 | + |
| 166 | + const result = triggerPayloadToChatRequest(longFilePathPayload) |
| 167 | + |
| 168 | + assert.strictEqual( |
| 169 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document |
| 170 | + ?.relativeFilePath?.length, |
| 171 | + filePathSizeLimit |
| 172 | + ) |
| 173 | + }) |
| 174 | + |
| 175 | + it('should preserve priority order', () => { |
| 176 | + const before1 = [5000, 30000, 40000, 20000, 15000, 25000] // Total: 135,000 |
| 177 | + const after1 = [5000, 30000, 40000, 20000, 5000, 0] // Total: 100,000 |
| 178 | + checkContextTruncationHelper(before1, after1) |
| 179 | + |
| 180 | + const before2 = [1000, 2000, 3000, 4000, 5000, 90000] // Total: 105,000 |
| 181 | + const after2 = [1000, 2000, 3000, 4000, 5000, 85000] // Total: 100,000 |
| 182 | + checkContextTruncationHelper(before2, after2) |
| 183 | + |
| 184 | + const before3 = [10000, 40000, 80000, 30000, 20000, 50000] // Total: 230,000 |
| 185 | + const after3 = [10000, 40000, 50000, 0, 0, 0] // Total: 100,000 |
| 186 | + checkContextTruncationHelper(before3, after3) |
| 187 | + |
| 188 | + const before4 = [5000, 5000, 150000, 5000, 5000, 5000] // Total: 175,000 |
| 189 | + const after4 = [5000, 5000, 90000, 0, 0, 0] // Total: 100,000 |
| 190 | + checkContextTruncationHelper(before4, after4) |
| 191 | + |
| 192 | + const before5 = [50000, 80000, 20000, 10000, 10000, 10000] // Total: 180,000 |
| 193 | + const after5 = [50000, 50000, 0, 0, 0, 0] // Total: 100,000 |
| 194 | + checkContextTruncationHelper(before5, after5) |
| 195 | + }) |
| 196 | + |
| 197 | + function checkContextTruncationHelper(before: number[], after: number[]) { |
| 198 | + const payload = createBaseTriggerPayload() |
| 199 | + const [userInputSize, promptSize, currentFileSize, ruleSize, fileSize, workspaceSize] = before |
| 200 | + |
| 201 | + payload.message = createLargeString(userInputSize, 'userInput-') |
| 202 | + payload.additionalContents = [createPrompt(promptSize), createRule(ruleSize), createFile(fileSize)] |
| 203 | + payload.fileText = createLargeString(currentFileSize, 'currentFile-') |
| 204 | + payload.relevantTextDocuments = [ |
| 205 | + { |
| 206 | + relativeFilePath: 'workspace.ts', |
| 207 | + text: createLargeString(workspaceSize, 'workspace-'), |
| 208 | + startLine: -1, |
| 209 | + endLine: -1, |
| 210 | + }, |
| 211 | + ] |
| 212 | + |
| 213 | + const result = triggerPayloadToChatRequest(payload) |
| 214 | + |
| 215 | + const userInputLength = result.conversationState.currentMessage?.userInputMessage?.content?.length |
| 216 | + const promptContext = |
| 217 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( |
| 218 | + (c) => c.name === 'prompt' |
| 219 | + )?.innerContext |
| 220 | + const currentFileLength = |
| 221 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document |
| 222 | + ?.text?.length |
| 223 | + const ruleContext = |
| 224 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( |
| 225 | + (c) => c.name === 'rule' |
| 226 | + )?.innerContext |
| 227 | + const fileContext = |
| 228 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( |
| 229 | + (c) => c.name === 'file' |
| 230 | + )?.innerContext |
| 231 | + const workspaceContext = |
| 232 | + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState |
| 233 | + ?.relevantDocuments?.[0]?.text |
| 234 | + |
| 235 | + // Verify priority ordering |
| 236 | + assert.strictEqual(userInputLength ?? 0, after[0]) |
| 237 | + assert.strictEqual(promptContext?.length ?? 0, after[1]) |
| 238 | + assert.strictEqual(currentFileLength ?? 0, after[2]) |
| 239 | + assert.strictEqual(ruleContext?.length ?? 0, after[3]) |
| 240 | + assert.strictEqual(fileContext?.length ?? 0, after[4]) |
| 241 | + assert.strictEqual(workspaceContext?.length ?? 0, after[5]) |
| 242 | + } |
| 243 | +}) |
0 commit comments