Skip to content

Commit be6fef8

Browse files
andrewyuqandrewyuq
committed
Make many fields in triggerPayload not undefined, so that we don't need many undefined checks becaus we are treating them as 0 in the end anyway
1 parent 27d93ac commit be6fef8

File tree

8 files changed

+406
-143
lines changed

8 files changed

+406
-143
lines changed

‎packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,17 +55,40 @@ export class InlineChatProvider {
5555
})
5656
return this.generateResponse(
5757
{
58-
message: message.message,
58+
message: message.message ?? '',
5959
trigger: ChatTriggerType.InlineChatMessage,
6060
query: message.message,
6161
codeSelection: context?.focusAreaContext?.selectionInsideExtendedCodeBlock,
62-
fileText: context?.focusAreaContext?.extendedCodeBlock,
62+
fileText: context?.focusAreaContext?.extendedCodeBlock ?? '',
6363
fileLanguage: context?.activeFileContext?.fileLanguage,
6464
filePath: context?.activeFileContext?.filePath,
6565
matchPolicy: context?.activeFileContext?.matchPolicy,
6666
codeQuery: context?.focusAreaContext?.names,
6767
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
6868
customization: getSelectedCustomization(),
69+
context: [],
70+
relevantTextDocuments: [],
71+
additionalContents: [],
72+
documentReferences: [],
73+
useRelevantDocuments: false,
74+
contextLengths: {
75+
additionalContextLengths: {
76+
fileContextLength: 0,
77+
promptContextLength: 0,
78+
ruleContextLength: 0,
79+
},
80+
truncatedAdditionalContextLengths: {
81+
fileContextLength: 0,
82+
promptContextLength: 0,
83+
ruleContextLength: 0,
84+
},
85+
workspaceContextLength: 0,
86+
truncatedWorkspaceContextLength: 0,
87+
userInputContextLength: 0,
88+
truncatedUserInputContextLength: 0,
89+
currentFileContextLength: 0,
90+
truncatedCurrentFileContextLength: 0,
91+
},
6992
},
7093
triggerID
7194
)
@@ -111,7 +134,10 @@ export class InlineChatProvider {
111134

112135
const request = triggerPayloadToChatRequest(triggerPayload)
113136
const session = this.sessionStorage.getSession(tabID)
114-
getLogger().info(`request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: %O`, request)
137+
getLogger().debug(
138+
`request from tab: ${tabID} conversationID: ${session.sessionIdentifier} request: %O`,
139+
request
140+
)
115141

116142
let response: GenerateAssistantResponseCommandOutput | undefined = undefined
117143
session.createNewTokenSource()

‎packages/core/src/amazonqTest/chat/controller/controller.ts‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,11 @@ import {
3636
import { UserIntent } from '@amzn/codewhisperer-streaming'
3737
import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil'
3838
import { createCodeWhispererChatStreamingClient } from '../../../shared/clients/codewhispererChatClient'
39-
import { ChatItemVotedMessage, ChatTriggerType } from '../../../codewhispererChat/controllers/chat/model'
39+
import {
40+
ChatItemVotedMessage,
41+
ChatTriggerType,
42+
TriggerPayload,
43+
} from '../../../codewhispererChat/controllers/chat/model'
4044
import { triggerPayloadToChatRequest } from '../../../codewhispererChat/controllers/chat/chatRequest/converter'
4145
import { EditorContentController } from '../../../amazonq/commons/controllers/contentController'
4246
import { amazonQTabSuffix } from '../../../shared/constants'
@@ -67,6 +71,7 @@ import { TargetFileInfo } from '../../../codewhisperer/client/codewhispereruserc
6771
import { submitFeedback } from '../../../feedback/vue/submitFeedback'
6872
import { placeholder } from '../../../shared/vscode/commands2'
6973
import { Auth } from '../../../auth/auth'
74+
import { defaultContextLengths } from '../../../codewhispererChat/constants'
7075

7176
export interface TestChatControllerEventEmitters {
7277
readonly tabOpened: vscode.EventEmitter<any>
@@ -916,7 +921,7 @@ export class TestController {
916921
// TODO: Write this entire gen response to basiccommands and call here.
917922
const editorText = await fs.readFileText(filePath)
918923

919-
const triggerPayload = {
924+
const triggerPayload: TriggerPayload = {
920925
query: `Generate unit tests for the following part of my code: ${message?.trim() || fileName}`,
921926
codeSelection: undefined,
922927
trigger: ChatTriggerType.ChatMessage,
@@ -928,6 +933,14 @@ export class TestController {
928933
codeQuery: undefined,
929934
userIntent: UserIntent.GENERATE_UNIT_TESTS,
930935
customization: getSelectedCustomization(),
936+
context: [],
937+
relevantTextDocuments: [],
938+
additionalContents: [],
939+
documentReferences: [],
940+
useRelevantDocuments: false,
941+
contextLengths: {
942+
...defaultContextLengths,
943+
},
931944
}
932945
const chatRequest = triggerPayloadToChatRequest(triggerPayload)
933946
const client = await createCodeWhispererChatStreamingClient()

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
*/
55
import * as path from 'path'
66
import fs from '../shared/fs/fs'
7+
import { ContextLengths } from './controllers/chat/model'
78

89
export const promptFileExtension = '.md'
910

@@ -20,3 +21,22 @@ export const getUserPromptsDirectory = () => {
2021
}
2122

2223
export const createSavedPromptCommandId = 'create-saved-prompt'
24+
25+
export const defaultContextLengths: ContextLengths = {
26+
additionalContextLengths: {
27+
fileContextLength: 0,
28+
promptContextLength: 0,
29+
ruleContextLength: 0,
30+
},
31+
truncatedAdditionalContextLengths: {
32+
fileContextLength: 0,
33+
promptContextLength: 0,
34+
ruleContextLength: 0,
35+
},
36+
workspaceContextLength: 0,
37+
truncatedWorkspaceContextLength: 0,
38+
userInputContextLength: 0,
39+
truncatedUserInputContextLength: 0,
40+
currentFileContextLength: 0,
41+
truncatedCurrentFileContextLength: 0,
42+
}

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

Lines changed: 16 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
import {
7-
ConversationState,
8-
CursorState,
9-
DocumentSymbol,
10-
RelevantTextDocument,
11-
SymbolType,
12-
TextDocument,
13-
} from '@amzn/codewhisperer-streaming'
6+
import { ConversationState, CursorState, DocumentSymbol, SymbolType, TextDocument } from '@amzn/codewhisperer-streaming'
147
import { AdditionalContentEntryAddition, ChatTriggerType, RelevantTextDocumentAddition, TriggerPayload } from '../model'
158
import { undefinedIfEmpty } from '../../../../shared/utilities/textUtilities'
169
import { getLogger } from '../../../../shared/logger/logger'
@@ -36,7 +29,6 @@ export const supportedLanguagesList = [
3629
]
3730

3831
const filePathSizeLimit = 4_000
39-
const customerMessageSizeLimit = 4_000
4032

4133
export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } {
4234
// Flexible truncation logic
@@ -81,7 +73,7 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
8173
)
8274

8375
getLogger().debug(
84-
`current request total payload size: ${userInputTruncationInfo.sizeAfter + userSpecificFilesTruncationInfo.sizeAfter + currentFileTruncationInfo.sizeAfter + userSpecificFilesTruncationInfo.sizeAfter + workspaceTruncationInfo.sizeAfter}`
76+
`current request total payload size: ${userInputTruncationInfo.sizeAfter + currentFileTruncationInfo.sizeAfter + userSpecificRulesTruncationInfo.sizeAfter + userSpecificFilesTruncationInfo.sizeAfter + workspaceTruncationInfo.sizeAfter}`
8577
)
8678

8779
// Filter out empty innerContext from additionalContents
@@ -92,11 +84,9 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
9284
}
9385

9486
// Filter out empty text from relevantTextDocuments
95-
if (triggerPayload.relevantTextDocuments !== undefined) {
96-
triggerPayload.relevantTextDocuments = triggerPayload.relevantTextDocuments.filter(
97-
(doc) => doc.text !== undefined && doc.text !== ''
98-
)
99-
}
87+
triggerPayload.relevantTextDocuments = triggerPayload.relevantTextDocuments.filter(
88+
(doc) => doc.text !== undefined && doc.text !== ''
89+
)
10090

10191
let document: TextDocument | undefined = undefined
10292
let cursorState: CursorState | undefined = undefined
@@ -154,10 +144,6 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
154144
}
155145
}
156146

157-
const relevantDocuments: RelevantTextDocument[] = triggerPayload.relevantTextDocuments
158-
? triggerPayload.relevantTextDocuments
159-
: []
160-
const useRelevantDocuments = triggerPayload.useRelevantDocuments
161147
// service will throw validation exception if string is empty
162148
const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn)
163149
const chatTriggerType = triggerPayload.trigger === ChatTriggerType.InlineChatMessage ? 'INLINE_CHAT' : 'MANUAL'
@@ -166,15 +152,13 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
166152
conversationState: {
167153
currentMessage: {
168154
userInputMessage: {
169-
content: triggerPayload.message
170-
? triggerPayload.message.substring(0, customerMessageSizeLimit)
171-
: '',
155+
content: triggerPayload.message,
172156
userInputMessageContext: {
173157
editorState: {
174158
document,
175159
cursorState,
176-
relevantDocuments,
177-
useRelevantDocuments,
160+
relevantDocuments: triggerPayload.relevantTextDocuments,
161+
useRelevantDocuments: triggerPayload.useRelevantDocuments,
178162
},
179163
additionalContext: triggerPayload.additionalContents,
180164
},
@@ -194,7 +178,7 @@ function preserveContexts(
194178
): FlexibleTruncationInfo {
195179
const typeToContextMap = new Map<
196180
ChatContextType,
197-
string | AdditionalContentEntryAddition[] | RelevantTextDocumentAddition[] | undefined
181+
string | AdditionalContentEntryAddition[] | RelevantTextDocumentAddition[]
198182
>([
199183
[ChatContextType.UserInput, triggerPayload.message],
200184
[ChatContextType.CurrentFile, triggerPayload.fileText],
@@ -212,45 +196,45 @@ function preserveContexts(
212196
}
213197

214198
const contexts = typeToContextMap.get(contextType)
215-
if (!contexts) {
216-
getLogger().debug(
217-
`Current request context size: type: ${contextType}, before: ${truncationInfo.sizeBefore}, after: ${truncationInfo.sizeAfter}`
218-
)
219-
return truncationInfo
220-
}
221-
222199
switch (contextType) {
223200
case ChatContextType.UserInput:
224201
truncationInfo = truncate(contexts as string, truncationInfo)
225202
triggerPayload.message = truncationInfo.textAfter
203+
triggerPayload.contextLengths.truncatedUserInputContextLength = truncationInfo.sizeAfter
226204
break
227205
case ChatContextType.CurrentFile:
228206
truncationInfo = truncate(contexts as string, truncationInfo)
229207
triggerPayload.fileText = truncationInfo.textAfter
208+
triggerPayload.contextLengths.truncatedCurrentFileContextLength = truncationInfo.sizeAfter
230209
break
231210
case ChatContextType.UserSpecificPrompts:
232211
truncationInfo = truncateUserSpecificContexts(
233212
contexts as AdditionalContentEntryAddition[],
234213
truncationInfo,
235214
'prompt'
236215
)
216+
triggerPayload.contextLengths.truncatedAdditionalContextLengths.promptContextLength =
217+
truncationInfo.sizeAfter
237218
break
238219
case ChatContextType.UserSpecificRules:
239220
truncationInfo = truncateUserSpecificContexts(
240221
contexts as AdditionalContentEntryAddition[],
241222
truncationInfo,
242223
'rule'
243224
)
225+
triggerPayload.contextLengths.truncatedAdditionalContextLengths.ruleContextLength = truncationInfo.sizeAfter
244226
break
245227
case ChatContextType.UserSpecificFiles:
246228
truncationInfo = truncateUserSpecificContexts(
247229
contexts as AdditionalContentEntryAddition[],
248230
truncationInfo,
249231
'file'
250232
)
233+
triggerPayload.contextLengths.truncatedAdditionalContextLengths.fileContextLength = truncationInfo.sizeAfter
251234
break
252235
case ChatContextType.Workspace:
253236
truncationInfo = truncateWorkspaceContexts(contexts as RelevantTextDocumentAddition[], truncationInfo)
237+
triggerPayload.contextLengths.truncatedWorkspaceContextLength = truncationInfo.sizeAfter
254238
break
255239
default:
256240
getLogger().warn(`Unexpected context type: ${contextType}`)

0 commit comments

Comments
 (0)