Skip to content

Commit 19973a8

Browse files
andrewyuqandrewyuq
committed
refactor: implement payload size management for chat requests
- Add size limits for different context types (user input, prompts, rules, files, workspace) - Implement truncation logic to preserve content around the middle of file text - Replace static contextMaxLength with workspaceChunkMaxSize for workspace content - Add debug logging for tracking payload sizes of different components - Filter out empty contexts from additionalContents and relevantTextDocuments - Update type definitions to include content type in AdditionalContentEntryAddition
1 parent bd37885 commit 19973a8

File tree

4 files changed

+143
-25
lines changed

4 files changed

+143
-25
lines changed

‎packages/core/src/codewhispererChat/constants.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,13 @@ import fs from '../shared/fs/fs'
77

88
export const promptFileExtension = '.md'
99

10+
// limit for each entry of @prompt, @rules, @files and @folder
1011
export const additionalContentInnerContextLimit = 8192
1112

1213
export const aditionalContentNameLimit = 1024
1314

14-
// temporary limit for @workspace and @file combined context length
15-
export const contextMaxLength = 40_000
15+
// limit for each chunk of @workspace
16+
export const workspaceChunkMaxSize = 40_960
1617

1718
export const getUserPromptsDirectory = () => {
1819
return path.join(fs.getUserHomeDir(), '.aws', 'amazonq', 'prompts')

‎packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
} from '@amzn/codewhisperer-streaming'
1414
import { ChatTriggerType, TriggerPayload } from '../model'
1515
import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
16+
import { getLogger } from '../../../../shared'
1617

1718
const fqnNameSizeDownLimit = 1
1819
const fqnNameSizeUpLimit = 256
@@ -38,6 +39,126 @@ const filePathSizeLimit = 4_000
3839
const customerMessageSizeLimit = 4_000
3940

4041
export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } {
42+
// truncate
43+
let remainingPayloadSize = 100_000
44+
// Type A context: Preserving userInput as much as possible
45+
if (triggerPayload.message !== undefined) {
46+
if (triggerPayload.message.length <= remainingPayloadSize) {
47+
remainingPayloadSize -= triggerPayload.message.length
48+
} else {
49+
triggerPayload.message = triggerPayload.message.substring(0, remainingPayloadSize)
50+
remainingPayloadSize = 0
51+
}
52+
}
53+
// TODO: send truncation telemetry if needed
54+
getLogger().debug(`current request user input size: ${triggerPayload.message?.length}`)
55+
56+
// Type B1(prompts) context: Preserving prompts as much as possible
57+
let totalPromptSize = 0
58+
if (triggerPayload.additionalContents !== undefined) {
59+
for (const additionalContent of triggerPayload.additionalContents) {
60+
if (additionalContent.type === 'prompt' && additionalContent.innerContext !== undefined) {
61+
if (additionalContent.innerContext.length <= remainingPayloadSize) {
62+
remainingPayloadSize -= additionalContent.innerContext.length
63+
} else {
64+
additionalContent.innerContext = additionalContent.innerContext.substring(0, remainingPayloadSize)
65+
remainingPayloadSize = 0
66+
}
67+
totalPromptSize += additionalContent.innerContext.length
68+
}
69+
}
70+
}
71+
72+
getLogger().debug(`current request total prompts size: ${totalPromptSize}`)
73+
74+
// Type C context: Preserving current file context as much as possible
75+
// truncate the text to keep texts in the middle instead of blindly truncating the tail
76+
if (triggerPayload.fileText !== undefined) {
77+
if (triggerPayload.fileText.length <= remainingPayloadSize) {
78+
remainingPayloadSize -= triggerPayload.fileText.length
79+
} else {
80+
// Calculate the middle point
81+
const middle = Math.floor(triggerPayload.fileText.length / 2)
82+
// Calculate how much text we can take from each side of the middle
83+
const halfRemaining = Math.floor(remainingPayloadSize / 2)
84+
// Get text from around the middle point
85+
const startPos = middle - halfRemaining
86+
const endPos = middle + halfRemaining
87+
88+
triggerPayload.fileText = triggerPayload.fileText.substring(startPos, endPos)
89+
remainingPayloadSize = 0
90+
}
91+
}
92+
getLogger().debug(`current request file content size: ${triggerPayload.fileText?.length}`)
93+
94+
// Type B1(rules) context: Preserving rules as much as possible
95+
let totalRulesSize = 0
96+
if (triggerPayload.additionalContents !== undefined) {
97+
for (const additionalContent of triggerPayload.additionalContents) {
98+
if (additionalContent.type === 'rule' && additionalContent.innerContext !== undefined) {
99+
if (additionalContent.innerContext.length <= remainingPayloadSize) {
100+
remainingPayloadSize -= additionalContent.innerContext.length
101+
} else {
102+
additionalContent.innerContext = additionalContent.innerContext.substring(0, remainingPayloadSize)
103+
remainingPayloadSize = 0
104+
}
105+
totalRulesSize += additionalContent.innerContext.length
106+
}
107+
}
108+
}
109+
110+
getLogger().debug(`current request rules size: ${totalRulesSize}`)
111+
112+
// Type B2(explicit @files) context: Preserving files as much as possible
113+
if (triggerPayload.additionalContents !== undefined) {
114+
for (const additionalContent of triggerPayload.additionalContents) {
115+
if (additionalContent.type === 'file' && additionalContent.innerContext !== undefined) {
116+
if (additionalContent.innerContext.length <= remainingPayloadSize) {
117+
remainingPayloadSize -= additionalContent.innerContext.length
118+
} else {
119+
additionalContent.innerContext = additionalContent.innerContext.substring(0, remainingPayloadSize)
120+
remainingPayloadSize = 0
121+
}
122+
}
123+
}
124+
}
125+
126+
// Type B3 @workspace context: Preserving workspace as much as possible
127+
let totalWorkspaceSize = 0
128+
if (triggerPayload.relevantTextDocuments !== undefined) {
129+
for (const relevantDocument of triggerPayload.relevantTextDocuments) {
130+
if (relevantDocument.text !== undefined) {
131+
if (relevantDocument.text.length <= remainingPayloadSize) {
132+
// remainingPayloadSize -= relevantDocument.text.length
133+
} else {
134+
// relevantDocument.text = relevantDocument.text.substring(0, remainingPayloadSize)
135+
// remainingPayloadSize = 0
136+
}
137+
totalWorkspaceSize += relevantDocument.text.length
138+
}
139+
}
140+
}
141+
142+
getLogger().debug(`current request workspace size: ${totalWorkspaceSize}`)
143+
144+
getLogger().debug(
145+
`current request total payload size: ${(triggerPayload.message?.length ?? 0) + totalPromptSize + (triggerPayload.fileText?.length ?? 0) + totalRulesSize + totalWorkspaceSize}`
146+
)
147+
148+
// Filter out empty innerContext from additionalContents
149+
if (triggerPayload.additionalContents !== undefined) {
150+
triggerPayload.additionalContents = triggerPayload.additionalContents.filter(
151+
(content) => content.innerContext !== undefined && content.innerContext !== ''
152+
)
153+
}
154+
155+
// Filter out empty text from relevantTextDocuments
156+
if (triggerPayload.relevantTextDocuments !== undefined) {
157+
triggerPayload.relevantTextDocuments = triggerPayload.relevantTextDocuments.filter(
158+
(doc) => doc.text !== undefined && doc.text !== ''
159+
)
160+
}
161+
41162
let document: TextDocument | undefined = undefined
42163
let cursorState: CursorState | undefined = undefined
43164

‎packages/core/src/codewhispererChat/controllers/chat/controller.ts‎

Lines changed: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ import {
7777
createSavedPromptCommandId,
7878
aditionalContentNameLimit,
7979
additionalContentInnerContextLimit,
80-
contextMaxLength,
80+
workspaceChunkMaxSize,
8181
} from '../../constants'
8282
import { ChatSession } from '../../clients/chat/v0/chat'
8383

@@ -967,7 +967,6 @@ export class ChatController {
967967
getLogger().verbose(`Could not get context command prompts: ${e}`)
968968
}
969969

970-
let currentContextLength = 0
971970
triggerPayload.additionalContents = []
972971
const emptyLengths = {
973972
fileContextLength: 0,
@@ -990,15 +989,16 @@ export class ChatController {
990989
name: prompt.name.substring(0, aditionalContentNameLimit),
991990
description: description.substring(0, aditionalContentNameLimit),
992991
innerContext: prompt.content.substring(0, additionalContentInnerContextLimit),
992+
type: contextType,
993993
}
994994
// make sure the relevantDocument + additionalContext
995995
// combined does not exceed 40k characters before generating the request payload.
996996
// Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation
997997
// TODO: Use a context length indicator
998-
if (currentContextLength + entry.innerContext.length > contextMaxLength) {
999-
getLogger().warn(`Selected context exceeds context size limit: ${entry.description} `)
1000-
break
1001-
}
998+
// if (currentContextLength + entry.innerContext.length > contextMaxLength) {
999+
// getLogger().warn(`Selected context exceeds context size limit: ${entry.description} `)
1000+
// break
1001+
// }
10021002

10031003
if (contextType === 'rule') {
10041004
triggerPayload.truncatedAdditionalContextLengths.ruleContextLength += entry.innerContext.length
@@ -1009,7 +1009,6 @@ export class ChatController {
10091009
}
10101010

10111011
triggerPayload.additionalContents.push(entry)
1012-
currentContextLength += entry.innerContext.length
10131012
let relativePath = path.relative(workspaceFolder, prompt.filePath)
10141013
// Handle user prompts outside the workspace
10151014
if (prompt.filePath.startsWith(getUserPromptsDirectory())) {
@@ -1067,25 +1066,20 @@ export class ChatController {
10671066
triggerPayload.message = triggerPayload.message.replace(/@workspace/, '')
10681067
if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
10691068
const start = performance.now()
1070-
let remainingContextLength = contextMaxLength
1071-
for (const additionalContent of triggerPayload.additionalContents || []) {
1072-
if (additionalContent.innerContext) {
1073-
remainingContextLength -= additionalContent.innerContext.length
1074-
}
1075-
}
1069+
// for (const additionalContent of triggerPayload.additionalContents || []) {
1070+
// if (additionalContent.innerContext) {
1071+
// remainingContextLength -= additionalContent.innerContext.length
1072+
// }
1073+
// }
10761074
triggerPayload.relevantTextDocuments = []
10771075
const relevantTextDocuments = await LspController.instance.query(triggerPayload.message)
10781076
for (const relevantDocument of relevantTextDocuments) {
10791077
if (relevantDocument.text !== undefined && relevantDocument.text.length > 0) {
1080-
if (remainingContextLength > relevantDocument.text.length) {
1081-
triggerPayload.relevantTextDocuments.push(relevantDocument)
1082-
remainingContextLength -= relevantDocument.text.length
1083-
} else {
1084-
getLogger().warn(
1085-
`Retrieved context exceeds context size limit: ${relevantDocument.relativeFilePath} `
1086-
)
1087-
break
1078+
if (relevantDocument.text.length > workspaceChunkMaxSize) {
1079+
relevantDocument.text = relevantDocument.text.substring(0, workspaceChunkMaxSize)
1080+
getLogger().debug(`Truncating @workspace chunk: ${relevantDocument.relativeFilePath} `)
10881081
}
1082+
triggerPayload.relevantTextDocuments.push(relevantDocument)
10891083
}
10901084
}
10911085

‎packages/core/src/codewhispererChat/controllers/chat/model.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ export interface TriggerPayload {
178178
readonly query: string | undefined
179179
readonly codeSelection: Selection | undefined
180180
readonly trigger: ChatTriggerType
181-
readonly fileText: string | undefined
181+
fileText: string | undefined
182182
readonly fileLanguage: string | undefined
183183
readonly filePath: string | undefined
184184
message: string | undefined
@@ -188,7 +188,7 @@ export interface TriggerPayload {
188188
readonly customization: Customization
189189
readonly context?: string[] | QuickActionCommand[]
190190
relevantTextDocuments?: RelevantTextDocumentAddition[]
191-
additionalContents?: AdditionalContentEntry[]
191+
additionalContents?: AdditionalContentEntryAddition[]
192192
// a reference to all documents used in chat payload
193193
// for providing better context transparency
194194
documentReferences?: DocumentReference[]
@@ -216,6 +216,8 @@ export type AdditionalContextInfo = {
216216
// TODO move this to API definition (or just use this across the codebase)
217217
export type RelevantTextDocumentAddition = RelevantTextDocument & { startLine: number; endLine: number }
218218

219+
export type AdditionalContentEntryAddition = AdditionalContentEntry & { type: string }
220+
219221
export interface DocumentReference {
220222
readonly relativeFilePath: string
221223
readonly lineRanges: Array<{ first: number; second: number }>

0 commit comments

Comments
 (0)