diff --git a/package-lock.json b/package-lock.json index f602ca6dbc0..49aac312dde 100644 --- a/package-lock.json +++ b/package-lock.json @@ -20,7 +20,7 @@ "vscode-nls-dev": "^4.0.4" }, "devDependencies": { - "@aws-toolkits/telemetry": "^1.0.308", + "@aws-toolkits/telemetry": "^1.0.311", "@playwright/browser-chromium": "^1.43.1", "@stylistic/eslint-plugin": "^2.11.0", "@types/he": "^1.2.3", @@ -10806,7 +10806,9 @@ } }, "node_modules/@aws-toolkits/telemetry": { - "version": "1.0.308", + "version": "1.0.311", + "resolved": "https://registry.npmjs.org/@aws-toolkits/telemetry/-/telemetry-1.0.311.tgz", + "integrity": "sha512-O8WhssqJLc6PPBwsONDPOmaZDGU56ZBRbRsyxw6faH/7aZEPOLQM6+ao9VhWeYfIkOy+yvcQiSc61zwM7DZ3OQ==", "dev": true, "license": "Apache-2.0", "dependencies": { diff --git a/package.json b/package.json index 1abf7ff4cf4..6364c093837 100644 --- a/package.json +++ b/package.json @@ -41,7 +41,7 @@ "skippedTestReport": "ts-node ./scripts/skippedTestReport.ts ./packages/amazonq/test/e2e/" }, "devDependencies": { - "@aws-toolkits/telemetry": "^1.0.308", + "@aws-toolkits/telemetry": "^1.0.311", "@playwright/browser-chromium": "^1.43.1", "@stylistic/eslint-plugin": "^2.11.0", "@types/he": "^1.2.3", diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts index aad41b3cd1b..af4ac95a036 100644 --- a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts +++ b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts @@ -55,17 +55,40 @@ export class InlineChatProvider { }) return this.generateResponse( { - message: message.message, + message: message.message ?? '', trigger: ChatTriggerType.InlineChatMessage, query: message.message, codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock, - fileText: context?.focusAreaContext?.extendedCodeBlock, + fileText: context?.focusAreaContext?.extendedCodeBlock ?? '', fileLanguage: context?.activeFileContext?.fileLanguage, filePath: context?.activeFileContext?.filePath, matchPolicy: context?.activeFileContext?.matchPolicy, codeQuery: context?.focusAreaContext?.names, userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message), customization: getSelectedCustomization(), + context: [], + relevantTextDocuments: [], + additionalContents: [], + documentReferences: [], + useRelevantDocuments: false, + contextLengths: { + additionalContextLengths: { + fileContextLength: 0, + promptContextLength: 0, + ruleContextLength: 0, + }, + truncatedAdditionalContextLengths: { + fileContextLength: 0, + promptContextLength: 0, + ruleContextLength: 0, + }, + workspaceContextLength: 0, + truncatedWorkspaceContextLength: 0, + userInputContextLength: 0, + truncatedUserInputContextLength: 0, + focusFileContextLength: 0, + truncatedFocusFileContextLength: 0, + }, }, triggerID ) @@ -111,7 +134,10 @@ export class InlineChatProvider { const request = triggerPayloadToChatRequest(triggerPayload) const session = this.sessionStorage.getSession(tabID) - getLogger().info(`request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: %O`, request) + getLogger().debug( + `request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: %O`, + request + ) let response: GenerateAssistantResponseCommandOutput | undefined = undefined session.createNewTokenSource() diff --git a/packages/amazonq/test/unit/codewhispererChat/controllers/chat/chatRequest/converter.test.ts b/packages/amazonq/test/unit/codewhispererChat/controllers/chat/chatRequest/converter.test.ts new file mode 100644 index 00000000000..d18b916229f --- /dev/null +++ b/packages/amazonq/test/unit/codewhispererChat/controllers/chat/chatRequest/converter.test.ts @@ -0,0 +1,243 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import { + ChatTriggerType, + filePathSizeLimit, + TriggerPayload, + triggerPayloadToChatRequest, +} from 'aws-core-vscode/codewhispererChat' + +describe('triggerPayloadToChatRequest', () => { + const mockBasicPayload: TriggerPayload = { + message: 'test message', + filePath: 'test/path.ts', + fileText: 'console.log("hello")', + fileLanguage: 'typescript', + additionalContents: [], + relevantTextDocuments: [], + useRelevantDocuments: false, + customization: { arn: 'test:arn' }, + trigger: ChatTriggerType.ChatMessage, + contextLengths: { + truncatedUserInputContextLength: 0, + truncatedFocusFileContextLength: 0, + truncatedWorkspaceContextLength: 0, + truncatedAdditionalContextLengths: { + promptContextLength: 0, + ruleContextLength: 0, + fileContextLength: 0, + }, + additionalContextLengths: { + fileContextLength: 0, + promptContextLength: 0, + ruleContextLength: 0, + }, + workspaceContextLength: 0, + userInputContextLength: 0, + focusFileContextLength: 0, + }, + context: [], + documentReferences: [], + query: undefined, + codeSelection: undefined, + matchPolicy: undefined, + codeQuery: undefined, + userIntent: undefined, + } + + const createLargeString = (size: number, prefix: string = '') => prefix + 'x'.repeat(size - prefix.length) + + const createPrompt = (size: number) => { + return { + name: 'prompt', + description: 'prompt', + relativePath: 'path-prompt', + type: 'prompt', + innerContext: createLargeString(size, 'prompt-'), + } + } + + const createRule = (size: number) => { + return { + name: 'rule', + description: 'rule', + relativePath: 'path-rule', + type: 'rule', + innerContext: createLargeString(size, 'rule-'), + } + } + + const createFile = (size: number) => { + return { + name: 'file', + description: 'file', + relativePath: 'path-file', + type: 'file', + innerContext: createLargeString(size, 'file-'), + } + } + + const createBaseTriggerPayload = (): TriggerPayload => ({ + ...mockBasicPayload, + message: '', + fileText: '', + filePath: 'test.ts', + fileLanguage: 'typescript', + customization: { arn: '' }, + }) + it('should convert basic trigger payload to chat request', () => { + const result = triggerPayloadToChatRequest(mockBasicPayload) + + assert.notEqual(result, undefined) + assert.strictEqual(result.conversationState.currentMessage?.userInputMessage?.content, 'test message') + assert.strictEqual(result.conversationState.chatTriggerType, 'MANUAL') + assert.strictEqual(result.conversationState.customizationArn, 'test:arn') + }) + + it('should handle empty file path', () => { + const emptyFilePathPayload = { + ...mockBasicPayload, + filePath: '', + } + + const result = triggerPayloadToChatRequest(emptyFilePathPayload) + + assert.strictEqual( + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document, + undefined + ) + }) + + it('should filter out empty additional contents', () => { + const payloadWithEmptyContents: TriggerPayload = { + ...mockBasicPayload, + additionalContents: [ + { name: 'prompt1', description: 'prompt1', relativePath: 'path1', type: 'prompt', innerContext: '' }, + { + name: 'prompt2', + description: 'prompt2', + relativePath: 'path2', + type: 'prompt', + innerContext: 'valid content', + }, + ], + } + + const result = triggerPayloadToChatRequest(payloadWithEmptyContents) + + assert.strictEqual( + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext + ?.length, + 1 + ) + assert.strictEqual( + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext.additionalContext?.[0] + .innerContext, + 'valid content' + ) + }) + + it('should handle unsupported programming language', () => { + const unsupportedLanguagePayload = { + ...mockBasicPayload, + fileLanguage: 'unsupported', + } + + const result = triggerPayloadToChatRequest(unsupportedLanguagePayload) + + assert.strictEqual( + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document + ?.programmingLanguage, + undefined + ) + }) + + it('should truncate file path if it exceeds limit', () => { + const longFilePath = 'a'.repeat(5000) + const longFilePathPayload = { + ...mockBasicPayload, + filePath: longFilePath, + } + + const result = triggerPayloadToChatRequest(longFilePathPayload) + + assert.strictEqual( + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document + ?.relativeFilePath?.length, + filePathSizeLimit + ) + }) + + it('should preserve priority order', () => { + const before1 = [5000, 30000, 40000, 20000, 15000, 25000] // Total: 135,000 + const after1 = [5000, 30000, 40000, 20000, 5000, 0] // Total: 100,000 + checkContextTruncationHelper(before1, after1) + + const before2 = [1000, 2000, 3000, 4000, 5000, 90000] // Total: 105,000 + const after2 = [1000, 2000, 3000, 4000, 5000, 85000] // Total: 100,000 + checkContextTruncationHelper(before2, after2) + + const before3 = [10000, 40000, 80000, 30000, 20000, 50000] // Total: 230,000 + const after3 = [10000, 40000, 50000, 0, 0, 0] // Total: 100,000 + checkContextTruncationHelper(before3, after3) + + const before4 = [5000, 5000, 150000, 5000, 5000, 5000] // Total: 175,000 + const after4 = [5000, 5000, 90000, 0, 0, 0] // Total: 100,000 + checkContextTruncationHelper(before4, after4) + + const before5 = [50000, 80000, 20000, 10000, 10000, 10000] // Total: 180,000 + const after5 = [50000, 50000, 0, 0, 0, 0] // Total: 100,000 + checkContextTruncationHelper(before5, after5) + }) + + function checkContextTruncationHelper(before: number[], after: number[]) { + const payload = createBaseTriggerPayload() + const [userInputSize, promptSize, currentFileSize, ruleSize, fileSize, workspaceSize] = before + + payload.message = createLargeString(userInputSize, 'userInput-') + payload.additionalContents = [createPrompt(promptSize), createRule(ruleSize), createFile(fileSize)] + payload.fileText = createLargeString(currentFileSize, 'currentFile-') + payload.relevantTextDocuments = [ + { + relativeFilePath: 'workspace.ts', + text: createLargeString(workspaceSize, 'workspace-'), + startLine: -1, + endLine: -1, + }, + ] + + const result = triggerPayloadToChatRequest(payload) + + const userInputLength = result.conversationState.currentMessage?.userInputMessage?.content?.length + const promptContext = + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( + (c) => c.name === 'prompt' + )?.innerContext + const currentFileLength = + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState?.document + ?.text?.length + const ruleContext = + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( + (c) => c.name === 'rule' + )?.innerContext + const fileContext = + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.additionalContext?.find( + (c) => c.name === 'file' + )?.innerContext + const workspaceContext = + result.conversationState.currentMessage?.userInputMessage?.userInputMessageContext?.editorState + ?.relevantDocuments?.[0]?.text + + // Verify priority ordering + assert.strictEqual(userInputLength ?? 0, after[0]) + assert.strictEqual(promptContext?.length ?? 0, after[1]) + assert.strictEqual(currentFileLength ?? 0, after[2]) + assert.strictEqual(ruleContext?.length ?? 0, after[3]) + assert.strictEqual(fileContext?.length ?? 0, after[4]) + assert.strictEqual(workspaceContext?.length ?? 0, after[5]) + } +}) diff --git a/packages/core/src/amazonqTest/chat/controller/controller.ts b/packages/core/src/amazonqTest/chat/controller/controller.ts index 4e29bb31464..cb5cbc2e851 100644 --- a/packages/core/src/amazonqTest/chat/controller/controller.ts +++ b/packages/core/src/amazonqTest/chat/controller/controller.ts @@ -36,7 +36,11 @@ import { import { UserIntent } from '@amzn/codewhisperer-streaming' import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil' import { createCodeWhispererChatStreamingClient } from '../../../shared/clients/codewhispererChatClient' -import { ChatItemVotedMessage, ChatTriggerType } from '../../../codewhispererChat/controllers/chat/model' +import { + ChatItemVotedMessage, + ChatTriggerType, + TriggerPayload, +} from '../../../codewhispererChat/controllers/chat/model' import { triggerPayloadToChatRequest } from '../../../codewhispererChat/controllers/chat/chatRequest/converter' import { EditorContentController } from '../../../amazonq/commons/controllers/contentController' import { amazonQTabSuffix } from '../../../shared/constants' @@ -67,6 +71,7 @@ import { TargetFileInfo } from '../../../codewhisperer/client/codewhispereruserc import { submitFeedback } from '../../../feedback/vue/submitFeedback' import { placeholder } from '../../../shared/vscode/commands2' import { Auth } from '../../../auth/auth' +import { defaultContextLengths } from '../../../codewhispererChat/constants' export interface TestChatControllerEventEmitters { readonly tabOpened: vscode.EventEmitter @@ -916,7 +921,7 @@ export class TestController { // TODO: Write this entire gen response to basiccommands and call here. const editorText = await fs.readFileText(filePath) - const triggerPayload = { + const triggerPayload: TriggerPayload = { query: `Generate unit tests for the following part of my code: ${message?.trim() || fileName}`, codeSelection: undefined, trigger: ChatTriggerType.ChatMessage, @@ -928,6 +933,14 @@ export class TestController { codeQuery: undefined, userIntent: UserIntent.GENERATE_UNIT_TESTS, customization: getSelectedCustomization(), + context: [], + relevantTextDocuments: [], + additionalContents: [], + documentReferences: [], + useRelevantDocuments: false, + contextLengths: { + ...defaultContextLengths, + }, } const chatRequest = triggerPayloadToChatRequest(triggerPayload) const client = await createCodeWhispererChatStreamingClient() diff --git a/packages/core/src/codewhispererChat/constants.ts b/packages/core/src/codewhispererChat/constants.ts index 4566d14ec64..84dd2dae292 100644 --- a/packages/core/src/codewhispererChat/constants.ts +++ b/packages/core/src/codewhispererChat/constants.ts @@ -4,18 +4,39 @@ */ import * as path from 'path' import fs from '../shared/fs/fs' +import { ContextLengths } from './controllers/chat/model' export const promptFileExtension = '.md' +// limit for each entry of @prompt, @rules, @files and @folder export const additionalContentInnerContextLimit = 8192 export const aditionalContentNameLimit = 1024 -// temporary limit for @workspace and @file combined context length -export const contextMaxLength = 40_000 +// limit for each chunk of @workspace +export const workspaceChunkMaxSize = 40_960 export const getUserPromptsDirectory = () => { return path.join(fs.getUserHomeDir(), '.aws', 'amazonq', 'prompts') } export const createSavedPromptCommandId = 'create-saved-prompt' + +export const defaultContextLengths: ContextLengths = { + additionalContextLengths: { + fileContextLength: 0, + promptContextLength: 0, + ruleContextLength: 0, + }, + truncatedAdditionalContextLengths: { + fileContextLength: 0, + promptContextLength: 0, + ruleContextLength: 0, + }, + workspaceContextLength: 0, + truncatedWorkspaceContextLength: 0, + userInputContextLength: 0, + truncatedUserInputContextLength: 0, + focusFileContextLength: 0, + truncatedFocusFileContextLength: 0, +} diff --git a/packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts b/packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts index 0a34463058e..be286122dc6 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts @@ -3,16 +3,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - ConversationState, - CursorState, - DocumentSymbol, - RelevantTextDocument, - SymbolType, - TextDocument, -} from '@amzn/codewhisperer-streaming' -import { ChatTriggerType, TriggerPayload } from '../model' +import { ConversationState, CursorState, DocumentSymbol, SymbolType, TextDocument } from '@amzn/codewhisperer-streaming' +import { AdditionalContentEntryAddition, ChatTriggerType, RelevantTextDocumentAddition, TriggerPayload } from '../model' import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities' +import { getLogger } from '../../../../shared/logger/logger' const fqnNameSizeDownLimit = 1 const fqnNameSizeUpLimit = 256 @@ -34,10 +28,66 @@ export const supportedLanguagesList = [ 'sql', ] -const filePathSizeLimit = 4_000 -const customerMessageSizeLimit = 4_000 +export const filePathSizeLimit = 4_000 export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } { + // Flexible truncation logic + const remainingPayloadSize = 100_000 + + // Type A context: Preserving user input as much as possible + const userInputTruncationInfo = preserveContexts(triggerPayload, remainingPayloadSize, ChatContextType.UserInput) + + // Type B1(prompts) context: Preserving @prompt as much as possible + const userSpecificPromptsTruncationInfo = preserveContexts( + triggerPayload, + userInputTruncationInfo.remainingPayloadSize, + ChatContextType.UserSpecificPrompts + ) + + // Type C context: Preserving current file context as much as possible + const currentFileTruncationInfo = preserveContexts( + triggerPayload, + userSpecificPromptsTruncationInfo.remainingPayloadSize, + ChatContextType.CurrentFile + ) + + // Type B1(rules) context: Preserving rules as much as possible + const userSpecificRulesTruncationInfo = preserveContexts( + triggerPayload, + currentFileTruncationInfo.remainingPayloadSize, + ChatContextType.UserSpecificRules + ) + + // Type B2(explicit @files) context: Preserving files as much as possible + const userSpecificFilesTruncationInfo = preserveContexts( + triggerPayload, + userSpecificRulesTruncationInfo.remainingPayloadSize, + ChatContextType.UserSpecificFiles + ) + + // Type B3 @workspace context: Preserving workspace as much as possible + const workspaceTruncationInfo = preserveContexts( + triggerPayload, + userSpecificFilesTruncationInfo.remainingPayloadSize, + ChatContextType.Workspace + ) + + getLogger().debug( + `current request total payload size: ${userInputTruncationInfo.sizeAfter + currentFileTruncationInfo.sizeAfter + userSpecificRulesTruncationInfo.sizeAfter + userSpecificFilesTruncationInfo.sizeAfter + workspaceTruncationInfo.sizeAfter}` + ) + + // Filter out empty innerContext from additionalContents + if (triggerPayload.additionalContents !== undefined) { + triggerPayload.additionalContents = triggerPayload.additionalContents.filter( + (content) => content.innerContext !== undefined && content.innerContext !== '' + ) + } + + // Filter out empty text from relevantTextDocuments + triggerPayload.relevantTextDocuments = triggerPayload.relevantTextDocuments.filter( + (doc) => doc.text !== undefined && doc.text !== '' + ) + let document: TextDocument | undefined = undefined let cursorState: CursorState | undefined = undefined @@ -94,10 +144,6 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c } } - const relevantDocuments: RelevantTextDocument[] = triggerPayload.relevantTextDocuments - ? triggerPayload.relevantTextDocuments - : [] - const useRelevantDocuments = triggerPayload.useRelevantDocuments // service will throw validation exception if string is empty const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn) const chatTriggerType = triggerPayload.trigger === ChatTriggerType.InlineChatMessage ? 'INLINE_CHAT' : 'MANUAL' @@ -106,15 +152,13 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c conversationState: { currentMessage: { userInputMessage: { - content: triggerPayload.message - ? triggerPayload.message.substring(0, customerMessageSizeLimit) - : '', + content: triggerPayload.message, userInputMessageContext: { editorState: { document, cursorState, - relevantDocuments, - useRelevantDocuments, + relevantDocuments: triggerPayload.relevantTextDocuments, + useRelevantDocuments: triggerPayload.useRelevantDocuments, }, additionalContext: triggerPayload.additionalContents, }, @@ -126,3 +170,151 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c }, } } + +function preserveContexts( + triggerPayload: TriggerPayload, + remainingPayloadSize: number, + contextType: ChatContextType +): FlexibleTruncationInfo { + const typeToContextMap = new Map< + ChatContextType, + string | AdditionalContentEntryAddition[] | RelevantTextDocumentAddition[] + >([ + [ChatContextType.UserInput, triggerPayload.message], + [ChatContextType.CurrentFile, triggerPayload.fileText], + [ChatContextType.UserSpecificPrompts, triggerPayload.additionalContents], + [ChatContextType.UserSpecificRules, triggerPayload.additionalContents], + [ChatContextType.UserSpecificFiles, triggerPayload.additionalContents], + [ChatContextType.Workspace, triggerPayload.relevantTextDocuments], + ]) + + let truncationInfo = { + remainingPayloadSize: remainingPayloadSize, + sizeBefore: 0, + sizeAfter: 0, + textAfter: '', + } + + const contexts = typeToContextMap.get(contextType) + switch (contextType) { + case ChatContextType.UserInput: + truncationInfo = truncate(contexts as string, truncationInfo) + triggerPayload.message = truncationInfo.textAfter + triggerPayload.contextLengths.truncatedUserInputContextLength = truncationInfo.sizeAfter + break + case ChatContextType.CurrentFile: + truncationInfo = truncate(contexts as string, truncationInfo) + triggerPayload.fileText = truncationInfo.textAfter + triggerPayload.contextLengths.truncatedFocusFileContextLength = truncationInfo.sizeAfter + break + case ChatContextType.UserSpecificPrompts: + truncationInfo = truncateUserSpecificContexts( + contexts as AdditionalContentEntryAddition[], + truncationInfo, + 'prompt' + ) + triggerPayload.contextLengths.truncatedAdditionalContextLengths.promptContextLength = + truncationInfo.sizeAfter + break + case ChatContextType.UserSpecificRules: + truncationInfo = truncateUserSpecificContexts( + contexts as AdditionalContentEntryAddition[], + truncationInfo, + 'rule' + ) + triggerPayload.contextLengths.truncatedAdditionalContextLengths.ruleContextLength = truncationInfo.sizeAfter + break + case ChatContextType.UserSpecificFiles: + truncationInfo = truncateUserSpecificContexts( + contexts as AdditionalContentEntryAddition[], + truncationInfo, + 'file' + ) + triggerPayload.contextLengths.truncatedAdditionalContextLengths.fileContextLength = truncationInfo.sizeAfter + break + case ChatContextType.Workspace: + truncationInfo = truncateWorkspaceContexts(contexts as RelevantTextDocumentAddition[], truncationInfo) + triggerPayload.contextLengths.truncatedWorkspaceContextLength = truncationInfo.sizeAfter + break + default: + getLogger().warn(`Unexpected context type: ${contextType}`) + return truncationInfo + } + + getLogger().debug( + `Current request context size: type: ${contextType}, before: ${truncationInfo.sizeBefore}, after: ${truncationInfo.sizeAfter}` + ) + return truncationInfo +} + +function truncateUserSpecificContexts( + contexts: AdditionalContentEntryAddition[], + truncationInfo: FlexibleTruncationInfo, + type: string +): FlexibleTruncationInfo { + for (const context of contexts) { + if (context.type !== type || !context.innerContext) { + continue + } + truncationInfo = truncate(context.innerContext, truncationInfo) + context.innerContext = truncationInfo.textAfter + } + return truncationInfo +} + +function truncateWorkspaceContexts( + contexts: RelevantTextDocumentAddition[], + truncationInfo: FlexibleTruncationInfo +): FlexibleTruncationInfo { + for (const context of contexts) { + if (!context.text) { + continue + } + truncationInfo = truncate(context.text, truncationInfo) + context.text = truncationInfo.textAfter + } + return truncationInfo +} + +function truncate( + textBefore: string, + truncationInfo: FlexibleTruncationInfo, + isCurrentFile: boolean = false +): FlexibleTruncationInfo { + const sizeBefore = truncationInfo.sizeBefore + textBefore.length + + // for all other types of contexts, we simply truncate the tail, + // for current file context, since it's expanded from the middle context, we truncate head and tail to preserve middle context + const middle = Math.floor(textBefore.length / 2) + const halfRemaining = Math.floor(truncationInfo.remainingPayloadSize / 2) + const startPos = isCurrentFile ? middle - halfRemaining : 0 + const endPos = isCurrentFile + ? middle + (truncationInfo.remainingPayloadSize - halfRemaining) + : Math.min(textBefore.length, truncationInfo.remainingPayloadSize) + const textAfter = textBefore.substring(startPos, endPos) + + const sizeAfter = truncationInfo.sizeAfter + textAfter.length + const remainingPayloadSize = truncationInfo.remainingPayloadSize - textAfter.length + return { + remainingPayloadSize, + sizeBefore, + sizeAfter, + textAfter, + } +} + +type FlexibleTruncationInfo = { + readonly remainingPayloadSize: number + readonly sizeBefore: number + readonly sizeAfter: number + readonly textAfter: string +} + +export enum ChatContextType { + UserInput = 'userInput', + CurrentFile = 'currentFile', + UserSpecificPrompts = 'userSpecificPrompts', + UserSpecificRules = 'userSpecificRules', + UserSpecificFiles = 'userSpecificFiles', + Workspace = 'workspace', +} diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index 846f3c6e445..8bcc36d835d 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -77,7 +77,8 @@ import { createSavedPromptCommandId, aditionalContentNameLimit, additionalContentInnerContextLimit, - contextMaxLength, + workspaceChunkMaxSize, + defaultContextLengths, } from '../../constants' import { ChatSession } from '../../clients/chat/v0/chat' @@ -720,13 +721,21 @@ export class ChatController { trigger: ChatTriggerType.ChatMessage, query: undefined, codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock, - fileText: context?.focusAreaContext?.extendedCodeBlock, + fileText: context?.focusAreaContext?.extendedCodeBlock ?? '', fileLanguage: context?.activeFileContext?.fileLanguage, filePath: context?.activeFileContext?.filePath, matchPolicy: context?.activeFileContext?.matchPolicy, codeQuery: context?.focusAreaContext?.names, userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command), customization: getSelectedCustomization(), + additionalContents: [], + relevantTextDocuments: [], + documentReferences: [], + useRelevantDocuments: false, + contextLengths: { + ...defaultContextLengths, + }, + context: [], }, triggerID ) @@ -793,17 +802,25 @@ export class ChatController { return this.generateResponse( { - message: message.message, + message: message.message ?? '', trigger: ChatTriggerType.ChatMessage, query: message.message, codeSelection: lastTriggerEvent.context?.focusAreaContext?.selectionInsideExtendedCodeBlock, - fileText: lastTriggerEvent.context?.focusAreaContext?.extendedCodeBlock, + fileText: lastTriggerEvent.context?.focusAreaContext?.extendedCodeBlock ?? '', fileLanguage: lastTriggerEvent.context?.activeFileContext?.fileLanguage, filePath: lastTriggerEvent.context?.activeFileContext?.filePath, matchPolicy: lastTriggerEvent.context?.activeFileContext?.matchPolicy, codeQuery: lastTriggerEvent.context?.focusAreaContext?.names, userIntent: message.userIntent, customization: getSelectedCustomization(), + contextLengths: { + ...defaultContextLengths, + }, + relevantTextDocuments: [], + additionalContents: [], + documentReferences: [], + useRelevantDocuments: false, + context: [], }, triggerID ) @@ -826,18 +843,25 @@ export class ChatController { }) return this.generateResponse( { - message: message.message, + message: message.message ?? '', trigger: ChatTriggerType.ChatMessage, query: message.message, codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock, - fileText: context?.focusAreaContext?.extendedCodeBlock, + fileText: context?.focusAreaContext?.extendedCodeBlock ?? '', fileLanguage: context?.activeFileContext?.fileLanguage, filePath: context?.activeFileContext?.filePath, matchPolicy: context?.activeFileContext?.matchPolicy, codeQuery: context?.focusAreaContext?.names, userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message), customization: getSelectedCustomization(), - context: message.context, + context: message.context ?? [], + relevantTextDocuments: [], + additionalContents: [], + documentReferences: [], + useRelevantDocuments: false, + contextLengths: { + ...defaultContextLengths, + }, }, triggerID ) @@ -907,12 +931,8 @@ export class ChatController { return rulesFiles } - private async resolveContextCommandPayload( - triggerPayload: TriggerPayload, - session: ChatSession - ): Promise { + private async resolveContextCommandPayload(triggerPayload: TriggerPayload, session: ChatSession) { const contextCommands: ContextCommandItem[] = [] - const relativePaths: string[] = [] // Check for workspace rules to add to context const workspaceRules = await this.collectWorkspaceRules() @@ -932,17 +952,16 @@ export class ChatController { triggerPayload.workspaceRulesCount = workspaceRules.length // Add context commands added by user to context - if (triggerPayload.context !== undefined && triggerPayload.context.length > 0) { - for (const context of triggerPayload.context) { - if (typeof context !== 'string' && context.route && context.route.length === 2) { - contextCommands.push({ - workspaceFolder: context.route[0] || '', - type: context.icon === 'folder' ? 'folder' : 'file', - relativePath: context.route[1] || '', - }) - } + for (const context of triggerPayload.context) { + if (typeof context !== 'string' && context.route && context.route.length === 2) { + contextCommands.push({ + workspaceFolder: context.route[0] || '', + type: context.icon === 'folder' ? 'folder' : 'file', + relativePath: context.route[1] || '', + }) } } + if (contextCommands.length === 0) { return [] } @@ -967,60 +986,31 @@ export class ChatController { getLogger().verbose(`Could not get context command prompts: ${e}`) } - let currentContextLength = 0 - triggerPayload.additionalContents = [] - const emptyLengths = { - fileContextLength: 0, - promptContextLength: 0, - ruleContextLength: 0, - } - triggerPayload.additionalContextLengths = emptyLengths - triggerPayload.truncatedAdditionalContextLengths = emptyLengths - - if (Array.isArray(prompts) && prompts.length > 0) { - triggerPayload.additionalContextLengths = this.telemetryHelper.getContextLengths(prompts) - for (const prompt of prompts.slice(0, 20)) { - // Add system prompt for user prompts and workspace rules - const contextType = this.telemetryHelper.getContextType(prompt) - const description = - contextType === 'rule' || contextType === 'prompt' - ? `You must follow the instructions in ${prompt.relativePath}. Below are lines ${prompt.startLine}-${prompt.endLine} of this file:\n` - : prompt.description - const entry = { - name: prompt.name.substring(0, aditionalContentNameLimit), - description: description.substring(0, aditionalContentNameLimit), - innerContext: prompt.content.substring(0, additionalContentInnerContextLimit), - } - // make sure the relevantDocument + additionalContext - // combined does not exceed 40k characters before generating the request payload. - // Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation - // TODO: Use a context length indicator - if (currentContextLength + entry.innerContext.length > contextMaxLength) { - getLogger().warn(`Selected context exceeds context size limit: ${entry.description} `) - break - } - - if (contextType === 'rule') { - triggerPayload.truncatedAdditionalContextLengths.ruleContextLength += entry.innerContext.length - } else if (contextType === 'prompt') { - triggerPayload.truncatedAdditionalContextLengths.promptContextLength += entry.innerContext.length - } else if (contextType === 'file') { - triggerPayload.truncatedAdditionalContextLengths.fileContextLength += entry.innerContext.length - } - - triggerPayload.additionalContents.push(entry) - currentContextLength += entry.innerContext.length - let relativePath = path.relative(workspaceFolder, prompt.filePath) - // Handle user prompts outside the workspace - if (prompt.filePath.startsWith(getUserPromptsDirectory())) { - relativePath = path.basename(prompt.filePath) - } - relativePaths.push(relativePath) + triggerPayload.contextLengths.additionalContextLengths = this.telemetryHelper.getContextLengths(prompts) + for (const prompt of prompts.slice(0, 20)) { + // Add system prompt for user prompts and workspace rules + const contextType = this.telemetryHelper.getContextType(prompt) + const description = + contextType === 'rule' || contextType === 'prompt' + ? `You must follow the instructions in ${prompt.relativePath}. Below are lines ${prompt.startLine}-${prompt.endLine} of this file:\n` + : prompt.description + + // Handle user prompts outside the workspace + const relativePath = prompt.filePath.startsWith(getUserPromptsDirectory()) + ? path.basename(prompt.filePath) + : path.relative(workspaceFolder, prompt.filePath) + + const entry = { + name: prompt.name.substring(0, aditionalContentNameLimit), + description: description.substring(0, aditionalContentNameLimit), + innerContext: prompt.content.substring(0, additionalContentInnerContextLimit), + type: contextType, + relativePath: relativePath, } + + triggerPayload.additionalContents.push(entry) } getLogger().info(`Retrieved chunks of additional context count: ${triggerPayload.additionalContents.length} `) - - return relativePaths } private async generateResponse( @@ -1057,35 +1047,23 @@ export class ChatController { } const session = this.sessionStorage.getSession(tabID) - const relativePathsOfContextCommandFiles = await this.resolveContextCommandPayload(triggerPayload, session) - triggerPayload.useRelevantDocuments = - triggerPayload.context?.some( - (context) => typeof context !== 'string' && context.command === '@workspace' - ) || false - triggerPayload.documentReferences = [] - if (triggerPayload.useRelevantDocuments && triggerPayload.message) { + await this.resolveContextCommandPayload(triggerPayload, session) + triggerPayload.useRelevantDocuments = triggerPayload.context.some( + (context) => typeof context !== 'string' && context.command === '@workspace' + ) + if (triggerPayload.useRelevantDocuments) { triggerPayload.message = triggerPayload.message.replace(/@workspace/, '') if (CodeWhispererSettings.instance.isLocalIndexEnabled()) { const start = performance.now() - let remainingContextLength = contextMaxLength - for (const additionalContent of triggerPayload.additionalContents || []) { - if (additionalContent.innerContext) { - remainingContextLength -= additionalContent.innerContext.length - } - } - triggerPayload.relevantTextDocuments = [] const relevantTextDocuments = await LspController.instance.query(triggerPayload.message) for (const relevantDocument of relevantTextDocuments) { - if (relevantDocument.text !== undefined && relevantDocument.text.length > 0) { - if (remainingContextLength > relevantDocument.text.length) { - triggerPayload.relevantTextDocuments.push(relevantDocument) - remainingContextLength -= relevantDocument.text.length - } else { - getLogger().warn( - `Retrieved context exceeds context size limit: ${relevantDocument.relativeFilePath} ` - ) - break + if (relevantDocument.text && relevantDocument.text.length > 0) { + triggerPayload.contextLengths.workspaceContextLength += relevantDocument.text.length + if (relevantDocument.text.length > workspaceChunkMaxSize) { + relevantDocument.text = relevantDocument.text.substring(0, workspaceChunkMaxSize) + getLogger().debug(`Truncating @workspace chunk: ${relevantDocument.relativeFilePath} `) } + triggerPayload.relevantTextDocuments.push(relevantDocument) } } @@ -1101,32 +1079,31 @@ export class ChatController { } } - triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments || []) - + triggerPayload.contextLengths.userInputContextLength = triggerPayload.message.length + triggerPayload.contextLengths.focusFileContextLength = triggerPayload.fileText.length const request = triggerPayloadToChatRequest(triggerPayload) + triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments) - if (triggerPayload.documentReferences !== undefined) { - const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map( - (doc) => doc.relativeFilePath - ) - const seen: string[] = [] - for (const relativePath of relativePathsOfContextCommandFiles) { - if (!relativePathsOfMergedRelevantDocuments.includes(relativePath) && !seen.includes(relativePath)) { - triggerPayload.documentReferences.push({ - relativeFilePath: relativePath, - lineRanges: [{ first: -1, second: -1 }], - }) - seen.push(relativePath) - } - } - if (triggerPayload.documentReferences) { - for (const doc of triggerPayload.documentReferences) { - session.contexts.set(doc.relativeFilePath, doc.lineRanges) - } + // Update context transparency after it's truncated dynamically to show users only the context sent. + const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map( + (doc) => doc.relativeFilePath + ) + const seen: string[] = [] + for (const additionalContent of triggerPayload.additionalContents) { + const relativePath = additionalContent.relativePath + if (!relativePathsOfMergedRelevantDocuments.includes(relativePath) && !seen.includes(relativePath)) { + triggerPayload.documentReferences.push({ + relativeFilePath: relativePath, + lineRanges: [{ first: -1, second: -1 }], + }) + seen.push(relativePath) } } + for (const doc of triggerPayload.documentReferences) { + session.contexts.set(doc.relativeFilePath, doc.lineRanges) + } - getLogger().info( + getLogger().debug( `request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: ${inspect(request, { depth: 12, })}` diff --git a/packages/core/src/codewhispererChat/controllers/chat/model.ts b/packages/core/src/codewhispererChat/controllers/chat/model.ts index ae0c6b61063..92878c515d9 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/model.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/model.ts @@ -178,27 +178,37 @@ export interface TriggerPayload { readonly query: string | undefined readonly codeSelection: Selection | undefined readonly trigger: ChatTriggerType - readonly fileText: string | undefined + fileText: string readonly fileLanguage: string | undefined readonly filePath: string | undefined - message: string | undefined + message: string readonly matchPolicy: MatchPolicy | undefined readonly codeQuery: CodeQuery | undefined readonly userIntent: UserIntent | undefined readonly customization: Customization - readonly context?: string[] | QuickActionCommand[] - relevantTextDocuments?: RelevantTextDocumentAddition[] - additionalContents?: AdditionalContentEntry[] + readonly context: string[] | QuickActionCommand[] + relevantTextDocuments: RelevantTextDocumentAddition[] + additionalContents: AdditionalContentEntryAddition[] // a reference to all documents used in chat payload // for providing better context transparency - documentReferences?: DocumentReference[] - useRelevantDocuments?: boolean + documentReferences: DocumentReference[] + useRelevantDocuments: boolean traceId?: string - additionalContextLengths?: AdditionalContextLengths - truncatedAdditionalContextLengths?: AdditionalContextLengths + contextLengths: ContextLengths workspaceRulesCount?: number } +export type ContextLengths = { + additionalContextLengths: AdditionalContextLengths + truncatedAdditionalContextLengths: AdditionalContextLengths + workspaceContextLength: number + truncatedWorkspaceContextLength: number + userInputContextLength: number + truncatedUserInputContextLength: number + focusFileContextLength: number + truncatedFocusFileContextLength: number +} + export type AdditionalContextLengths = { fileContextLength: number promptContextLength: number @@ -216,6 +226,8 @@ export type AdditionalContextInfo = { // TODO move this to API definition (or just use this across the codebase) export type RelevantTextDocumentAddition = RelevantTextDocument & { startLine: number; endLine: number } +export type AdditionalContentEntryAddition = AdditionalContentEntry & { type: string; relativePath: string } + export interface DocumentReference { readonly relativeFilePath: string readonly lineRanges: Array<{ first: number; second: number }> diff --git a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts index eb99b8c5dca..5d5cc09056d 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts @@ -506,7 +506,7 @@ export class CWCTelemetryHelper { cwsprChatUserIntent: this.getUserIntentForTelemetry(triggerPayload.userIntent), cwsprChatHasCodeSnippet: triggerPayload.codeSelection && !triggerPayload.codeSelection.isEmpty, cwsprChatProgrammingLanguage: triggerPayload.fileLanguage, - cwsprChatActiveEditorTotalCharacters: triggerPayload.fileText?.length, + cwsprChatActiveEditorTotalCharacters: triggerPayload.fileText.length, cwsprChatActiveEditorImportCount: triggerPayload.codeQuery?.fullyQualifiedNames?.used?.length, cwsprChatResponseCodeSnippetCount: message.totalNumberOfCodeBlocksInResponse, cwsprChatResponseCode: message.responseCode, @@ -523,26 +523,32 @@ export class CWCTelemetryHelper { cwsprChatFullServerResponseLatency: this.conversationStreamTotalTime.get(message.tabID) ?? 0, cwsprChatTimeBetweenDisplays: JSON.stringify(this.getTimeBetweenChunks(tabID, this.displayTimeForChunks)), cwsprChatFullDisplayLatency: fullDisplayLatency, - cwsprChatRequestLength: triggerPayload.message?.length ?? 0, + cwsprChatRequestLength: triggerPayload.message.length, cwsprChatResponseLength: message.messageLength, cwsprChatConversationType: 'Chat', credentialStartUrl: AuthUtil.instance.startUrl, codewhispererCustomizationArn: triggerPayload.customization.arn, cwsprChatHasProjectContext: hasProjectLevelContext, - cwsprChatHasContextList: triggerPayload.documentReferences && triggerPayload.documentReferences?.length > 0, + cwsprChatHasContextList: triggerPayload.documentReferences.length > 0, cwsprChatFolderContextCount: contextCounts.folderContextCount, cwsprChatFileContextCount: contextCounts.fileContextCount, - cwsprChatFileContextLength: triggerPayload.additionalContextLengths?.fileContextLength ?? 0, + cwsprChatFileContextLength: triggerPayload.contextLengths.additionalContextLengths.fileContextLength, cwsprChatFileContextTruncatedLength: - triggerPayload.truncatedAdditionalContextLengths?.fileContextLength ?? 0, + triggerPayload.contextLengths.truncatedAdditionalContextLengths.fileContextLength, cwsprChatRuleContextCount: triggerPayload.workspaceRulesCount, - cwsprChatRuleContextLength: triggerPayload.additionalContextLengths?.ruleContextLength ?? 0, + cwsprChatRuleContextLength: triggerPayload.contextLengths.additionalContextLengths.ruleContextLength, cwsprChatRuleContextTruncatedLength: - triggerPayload.truncatedAdditionalContextLengths?.ruleContextLength ?? 0, + triggerPayload.contextLengths.truncatedAdditionalContextLengths.ruleContextLength, cwsprChatPromptContextCount: contextCounts.promptContextCount, - cwsprChatPromptContextLength: triggerPayload.additionalContextLengths?.promptContextLength ?? 0, + cwsprChatPromptContextLength: triggerPayload.contextLengths.additionalContextLengths.promptContextLength, cwsprChatPromptContextTruncatedLength: - triggerPayload.truncatedAdditionalContextLengths?.promptContextLength ?? 0, + triggerPayload.contextLengths.truncatedAdditionalContextLengths.promptContextLength, + cwsprChatFocusFileContextLength: triggerPayload.contextLengths.focusFileContextLength, + cwsprChatFocusFileContextTruncatedLength: triggerPayload.contextLengths.truncatedFocusFileContextLength, + cwsprChatUserInputContextLength: triggerPayload.contextLengths.userInputContextLength, + cwsprChatUserInputContextTruncatedLength: triggerPayload.contextLengths.truncatedUserInputContextLength, + cwsprChatWorkspaceContextLength: triggerPayload.contextLengths.workspaceContextLength, + cwsprChatWorkspaceContextTruncatedLength: triggerPayload.contextLengths.truncatedWorkspaceContextLength, traceId, } diff --git a/packages/core/src/codewhispererChat/index.ts b/packages/core/src/codewhispererChat/index.ts index e473203caf5..b47115fbc4a 100644 --- a/packages/core/src/codewhispererChat/index.ts +++ b/packages/core/src/codewhispererChat/index.ts @@ -7,7 +7,7 @@ export { FocusAreaContextExtractor } from './editor/context/focusArea/focusAreaE export { TryChatCodeLensProvider, resolveModifierKey, tryChatCodeLensCommand } from './editor/codelens' export { focusAmazonQPanel } from './commands/registerCommands' export { ChatSession } from './clients/chat/v0/chat' -export { triggerPayloadToChatRequest } from './controllers/chat/chatRequest/converter' +export { triggerPayloadToChatRequest, filePathSizeLimit } from './controllers/chat/chatRequest/converter' export { ChatTriggerType, PromptMessage, TriggerPayload } from './controllers/chat/model' export { UserIntentRecognizer } from './controllers/chat/userIntent/userIntentRecognizer' export { EditorContextExtractor } from './editor/context/extractor'