Skip to content

Commit 77643cf

Browse files
ashishrp-awsashishrp-aws
committed
Added chat history storage to the Q chat for agentic chat
1 parent 586f69a commit 77643cf

File tree

14 files changed

+1522
-123
lines changed

14 files changed

+1522
-123
lines changed

‎packages/core/src/codewhisperer/client/user-service-2.json

Lines changed: 682 additions & 35 deletions
Large diffs are not rendered by default.

‎packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,6 @@ export class ChatSession {
6060
async chatSso(chatRequest: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
6161
const client = await createCodeWhispererChatStreamingClient()
6262

63-
if (this.sessionId !== undefined && chatRequest.conversationState !== undefined) {
64-
chatRequest.conversationState.conversationId = this.sessionId
65-
}
66-
6763
const response = await client.generateAssistantResponse(chatRequest)
6864
if (!response.generateAssistantResponseResponse) {
6965
throw new ToolkitError(

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,10 +119,12 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { c
119119
additionalContext: triggerPayload.additionalContents,
120120
},
121121
userIntent: triggerPayload.userIntent,
122+
origin: 'IDE',
122123
},
123124
},
124125
chatTriggerType,
125126
customizationArn: customizationArn,
127+
history: triggerPayload.chatHistory,
126128
},
127129
}
128130
}

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ import {
8080
contextMaxLength,
8181
} from '../../constants'
8282
import { ChatSession } from '../../clients/chat/v0/chat'
83+
import { ChatHistoryManager } from '../../storages/chatHistory'
8384

8485
export interface ChatControllerMessagePublishers {
8586
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -141,6 +142,7 @@ export class ChatController {
141142
private readonly userIntentRecognizer: UserIntentRecognizer
142143
private readonly telemetryHelper: CWCTelemetryHelper
143144
private userPromptsWatcher: vscode.FileSystemWatcher | undefined
145+
private readonly chatHistoryManager: ChatHistoryManager
144146

145147
public constructor(
146148
private readonly chatControllerMessageListeners: ChatControllerMessageListeners,
@@ -158,6 +160,7 @@ export class ChatController {
158160
this.editorContentController = new EditorContentController()
159161
this.promptGenerator = new PromptsGenerator()
160162
this.userIntentRecognizer = new UserIntentRecognizer()
163+
this.chatHistoryManager = new ChatHistoryManager()
161164

162165
onDidChangeAmazonQVisibility((visible) => {
163166
if (visible) {
@@ -395,6 +398,7 @@ export class ChatController {
395398

396399
private async processTabCloseMessage(message: TabClosedMessage) {
397400
this.sessionStorage.deleteSession(message.tabID)
401+
this.chatHistoryManager.clear()
398402
this.triggerEventsStorage.removeTabEvents(message.tabID)
399403
// this.telemetryHelper.recordCloseChat(message.tabID)
400404
}
@@ -654,6 +658,7 @@ export class ChatController {
654658
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
655659

656660
this.sessionStorage.deleteSession(tabID)
661+
this.chatHistoryManager.clear()
657662
}
658663

659664
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -714,6 +719,13 @@ export class ChatController {
714719
command,
715720
})
716721

722+
this.chatHistoryManager.appendUserMessage({
723+
userInputMessage: {
724+
content: prompt,
725+
userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
726+
},
727+
})
728+
717729
return this.generateResponse(
718730
{
719731
message: prompt,
@@ -727,6 +739,7 @@ export class ChatController {
727739
codeQuery: context?.focusAreaContext?.names,
728740
userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
729741
customization: getSelectedCustomization(),
742+
chatHistory: this.chatHistoryManager.getHistory(),
730743
},
731744
triggerID
732745
)
@@ -766,6 +779,7 @@ export class ChatController {
766779
switch (message.command) {
767780
case 'clear':
768781
this.sessionStorage.deleteSession(message.tabID)
782+
this.chatHistoryManager.clear()
769783
this.triggerEventsStorage.removeTabEvents(message.tabID)
770784
recordTelemetryChatRunCommand('clear')
771785
return
@@ -791,6 +805,13 @@ export class ChatController {
791805
context: lastTriggerEvent.context,
792806
})
793807

808+
this.chatHistoryManager.appendUserMessage({
809+
userInputMessage: {
810+
content: message.message,
811+
userIntent: message.userIntent,
812+
},
813+
})
814+
794815
return this.generateResponse(
795816
{
796817
message: message.message,
@@ -804,6 +825,7 @@ export class ChatController {
804825
codeQuery: lastTriggerEvent.context?.focusAreaContext?.names,
805826
userIntent: message.userIntent,
806827
customization: getSelectedCustomization(),
828+
chatHistory: this.chatHistoryManager.getHistory(),
807829
},
808830
triggerID
809831
)
@@ -824,6 +846,12 @@ export class ChatController {
824846
type: 'chat_message',
825847
context,
826848
})
849+
this.chatHistoryManager.appendUserMessage({
850+
userInputMessage: {
851+
content: message.message,
852+
userIntent: message.userIntent,
853+
},
854+
})
827855
return this.generateResponse(
828856
{
829857
message: message.message,
@@ -838,6 +866,7 @@ export class ChatController {
838866
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
839867
customization: getSelectedCustomization(),
840868
context: message.context,
869+
chatHistory: this.chatHistoryManager.getHistory(),
841870
},
842871
triggerID
843872
)
@@ -1104,6 +1133,15 @@ export class ChatController {
11041133
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments || [])
11051134

11061135
const request = triggerPayloadToChatRequest(triggerPayload)
1136+
if (
1137+
this.chatHistoryManager.getConversationId() !== undefined &&
1138+
this.chatHistoryManager.getConversationId() !== ''
1139+
) {
1140+
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1141+
} else {
1142+
this.chatHistoryManager.setConversationId(randomUUID())
1143+
request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
1144+
}
11071145

11081146
if (triggerPayload.documentReferences !== undefined) {
11091147
const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map(
@@ -1157,7 +1195,14 @@ export class ChatController {
11571195
response.$metadata.requestId
11581196
} metadata: ${inspect(response.$metadata, { depth: 12 })}`
11591197
)
1160-
await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload)
1198+
await this.messenger.sendAIResponse(
1199+
response,
1200+
session,
1201+
tabID,
1202+
triggerID,
1203+
triggerPayload,
1204+
this.chatHistoryManager
1205+
)
11611206
} catch (e: any) {
11621207
this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
11631208
// clears session, record telemetry before this call

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import { extractCodeBlockLanguage } from '../../../../shared/markdown'
3838
import { extractAuthFollowUp } from '../../../../amazonq/util/authUtils'
3939
import { helpMessage } from '../../../../amazonq/webview/ui/texts/constants'
4040
import { ChatItemButton, ChatItemFormItem, MynahUIDataModel } from '@aws/mynah-ui'
41+
import { ChatHistoryManager } from '../../../storages/chatHistory'
4142

4243
export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
4344

@@ -121,7 +122,8 @@ export class Messenger {
121122
session: ChatSession,
122123
tabID: string,
123124
triggerID: string,
124-
triggerPayload: TriggerPayload
125+
triggerPayload: TriggerPayload,
126+
chatHistoryManager: ChatHistoryManager
125127
) {
126128
let message = ''
127129
const messageID = response.$metadata.requestId ?? ''
@@ -331,6 +333,15 @@ export class Messenger {
331333
)
332334
)
333335

336+
chatHistoryManager.pushAssistantMessage({
337+
assistantResponseMessage: {
338+
messageId: messageID,
339+
content: message,
340+
references: codeReference,
341+
// TODO: Add tools data and follow up prompt details
342+
},
343+
})
344+
334345
getLogger().info(
335346
`All events received. requestId=%s counts=%s`,
336347
response.$metadata.requestId,

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*/
55

66
import * as vscode from 'vscode'
7-
import { AdditionalContentEntry, RelevantTextDocument, UserIntent } from '@amzn/codewhisperer-streaming'
7+
import { AdditionalContentEntry, ChatMessage, RelevantTextDocument, UserIntent } from '@amzn/codewhisperer-streaming'
88
import { MatchPolicy, CodeQuery } from '../../clients/chat/v0/model'
99
import { Selection } from 'vscode'
1010
import { TabOpenType } from '../../../amazonq/webview/ui/storages/tabsStorage'
@@ -197,6 +197,7 @@ export interface TriggerPayload {
197197
additionalContextLengths?: AdditionalContextLengths
198198
truncatedAdditionalContextLengths?: AdditionalContextLengths
199199
workspaceRulesCount?: number
200+
chatHistory?: ChatMessage[]
200201
}
201202

202203
export type AdditionalContextLengths = {

‎packages/core/src/codewhispererChat/storages/chatHistory.ts

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
import { ChatMessage } from '@amzn/codewhisperer-streaming'
6+
import { randomUUID } from '../../shared/crypto'
7+
import { getLogger } from '../../shared/logger/logger'
8+
9+
// Maximum number of messages to keep in history
10+
const MaxConversationHistoryLength = 100
11+
12+
/**
13+
* ChatHistoryManager handles the storage and manipulation of chat history
14+
* for CodeWhisperer Chat sessions.
15+
*/
16+
export class ChatHistoryManager {
17+
private conversationId: string
18+
private history: ChatMessage[] = []
19+
private logger = getLogger()
20+
private lastUserMessage?: ChatMessage
21+
22+
constructor() {
23+
this.conversationId = randomUUID()
24+
this.logger.info(`Generated new conversation id: ${this.conversationId}`)
25+
}
26+
27+
/**
28+
* Get the conversation ID
29+
*/
30+
public getConversationId(): string {
31+
return this.conversationId
32+
}
33+
34+
public setConversationId(conversationId: string) {
35+
this.conversationId = conversationId
36+
}
37+
38+
/**
39+
* Get the full chat history
40+
*/
41+
public getHistory(): ChatMessage[] {
42+
return [...this.history]
43+
}
44+
45+
/**
46+
* Clear the conversation history
47+
*/
48+
public clear(): void {
49+
this.history = []
50+
this.conversationId = ''
51+
}
52+
53+
/**
54+
* Append a new user message to be sent
55+
*/
56+
public appendUserMessage(newMessage: ChatMessage): void {
57+
this.fixHistory()
58+
if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
59+
this.logger.warn('input must not be empty when adding new messages')
60+
// const emptyMessage: ChatMessage = {
61+
// ...newMessage,
62+
// userInputMessage: {
63+
// ...newMessage.userInputMessage,
64+
// content: 'Empty user input',
65+
// },
66+
// }
67+
// this.history.push(emptyMessage)
68+
}
69+
this.lastUserMessage = newMessage
70+
this.history.push(newMessage)
71+
}
72+
73+
/**
74+
* Push an assistant message to the history
75+
*/
76+
public pushAssistantMessage(newMessage: ChatMessage): void {
77+
if (newMessage !== undefined && this.lastUserMessage !== undefined) {
78+
this.logger.warn('last Message should not be defined when pushing an assistant message')
79+
}
80+
this.history.push(newMessage)
81+
}
82+
83+
/**
84+
* Fixes the history to maintain the following invariants:
85+
* 1. The history length is <= MAX_CONVERSATION_HISTORY_LENGTH. Oldest messages are dropped.
86+
* 2. The first message is from the user. Oldest messages are dropped if needed.
87+
* 3. The last message is from the assistant. The last message is dropped if it is from the user.
88+
* 4. If the last message is from the assistant and it contains tool uses, and a next user
89+
* message is set without tool results, then the user message will have cancelled tool results.
90+
*/
91+
public fixHistory(): void {
92+
// Trim the conversation history if it exceeds the maximum length
93+
if (this.history.length > MaxConversationHistoryLength) {
94+
// Find the second oldest user message to be the new starting point
95+
const secondUserMessageIndex = this.history
96+
.slice(1) // Skip the first message which might be from the user
97+
.findIndex((msg) => !msg.userInputMessage?.content || msg.userInputMessage?.content.trim() === '')
98+
99+
if (secondUserMessageIndex !== -1) {
100+
// +1 because we sliced off the first element
101+
this.logger.debug(`Removing the first ${secondUserMessageIndex + 1} elements in the history`)
102+
this.history = this.history.slice(secondUserMessageIndex + 1)
103+
} else {
104+
this.logger.debug('No valid starting user message found in the history, clearing')
105+
this.history = []
106+
}
107+
}
108+
109+
// Ensure the last message is from the assistant
110+
111+
if (this.history.length > 0 && this.history[this.history.length - 1].userInputMessage !== undefined) {
112+
this.logger.debug('Last message in history is from the user, dropping')
113+
this.history.pop()
114+
}
115+
116+
// TODO: If the last message from the assistant contains tool uses, ensure the next user message contains tool results
117+
}
118+
}

0 commit comments

Comments
 (0)