Skip to content

Commit b5cd5c3

Browse files
authored
feat(amazonq): Use Chat History DB instead of in-memory store (aws#6985)
## Problem We need to switch to using the Chat History DB instead of the local store ## Solution - Add the same history fixing functionality to Chat History DB as the former chatHistory local store - Replace all uses of the local store - Maintain a sessionId/ConversationId for each tab ## Tested manually --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent b222e0b commit b5cd5c3

File tree

6 files changed

+251
-370
lines changed

6 files changed

+251
-370
lines changed

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

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,15 @@ import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDev
1616
import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker'
1717
import { DocumentReference, PromptMessage } from '../../../controllers/chat/model'
1818
import { FsWriteBackup } from '../../../../codewhispererChat/tools/fsWrite'
19+
import { randomUUID } from '../../../../shared/crypto'
1920

2021
export type ToolUseWithError = {
2122
toolUse: ToolUse
2223
error: Error | undefined
2324
}
24-
import { getLogger } from '../../../../shared/logger/logger'
25-
import { randomUUID } from '../../../../shared/crypto'
2625

2726
export class ChatSession {
28-
private sessionId?: string
27+
private sessionId: string
2928
/**
3029
* _readFiles = list of files read from the project to gather context before generating response.
3130
* _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user
@@ -50,7 +49,7 @@ export class ChatSession {
5049
// TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
5150
// e.g. root_a/file1 vs root_b/file1
5251
relativePathToWorkspaceRoot: Map<string, string> = new Map()
53-
public get sessionIdentifier(): string | undefined {
52+
public get sessionIdentifier(): string {
5453
return this.sessionId
5554
}
5655
public get messageIdToUpdate(): string | undefined {
@@ -105,13 +104,14 @@ export class ChatSession {
105104

106105
constructor() {
107106
this.createNewTokenSource()
107+
this.sessionId = randomUUID()
108108
}
109109

110110
createNewTokenSource() {
111111
this.tokenSource = new vscode.CancellationTokenSource()
112112
}
113113

114-
public setSessionID(id?: string) {
114+
public setSessionID(id: string) {
115115
this.sessionId = id
116116
}
117117
public get readFiles(): DocumentReference[] {
@@ -148,14 +148,6 @@ export class ChatSession {
148148
)
149149
}
150150

151-
const responseStream = response.sendMessageResponse
152-
for await (const event of responseStream) {
153-
if ('messageMetadataEvent' in event) {
154-
this.sessionId = event.messageMetadataEvent?.conversationId
155-
break
156-
}
157-
}
158-
159151
UserWrittenCodeTracker.instance.onQFeatureInvoked()
160152
return response
161153
}
@@ -170,12 +162,6 @@ export class ChatSession {
170162
)
171163
}
172164

173-
this.sessionId = response.conversationId
174-
if (this.sessionId?.length === 0) {
175-
getLogger().debug(`Session ID: ${this.sessionId} is empty. Generating random UUID`)
176-
this.sessionId = randomUUID()
177-
}
178-
179165
UserWrittenCodeTracker.instance.onQFeatureInvoked()
180166

181167
return response

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

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -97,10 +97,10 @@ import { amazonQTabSuffix } from '../../../shared/constants'
9797
import { OutputKind } from '../../tools/toolShared'
9898
import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils'
9999
import { ChatStream } from '../../tools/chatStream'
100-
import { ChatHistoryStorage } from '../../storages/chatHistoryStorage'
101100
import { tempDirPath } from '../../../shared/filesystemUtilities'
102101
import { Database } from '../../../shared/db/chatDb/chatDb'
103102
import { TabBarController } from './tabBarController'
103+
import { messageToChatMessage } from '../../../shared/db/chatDb/util'
104104

105105
export interface ChatControllerMessagePublishers {
106106
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -175,7 +175,6 @@ export class ChatController {
175175
private readonly userIntentRecognizer: UserIntentRecognizer
176176
private readonly telemetryHelper: CWCTelemetryHelper
177177
private userPromptsWatcher: vscode.FileSystemWatcher | undefined
178-
private readonly chatHistoryStorage: ChatHistoryStorage
179178
private chatHistoryDb = Database.getInstance()
180179
private cancelTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource()
181180

@@ -195,7 +194,6 @@ export class ChatController {
195194
this.editorContentController = new EditorContentController()
196195
this.promptGenerator = new PromptsGenerator()
197196
this.userIntentRecognizer = new UserIntentRecognizer()
198-
this.chatHistoryStorage = new ChatHistoryStorage()
199197
this.tabBarController = new TabBarController(this.messenger)
200198

201199
onDidChangeAmazonQVisibility((visible) => {
@@ -417,7 +415,7 @@ export class ChatController {
417415
const session = this.sessionStorage.getSession(message.tabID)
418416
session.tokenSource.cancel()
419417
this.messenger.sendEmptyMessage(message.tabID, '', undefined)
420-
this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
418+
this.chatHistoryDb.clearRecentHistory(message.tabID)
421419
this.telemetryHelper.recordInteractionWithAgenticChat(AgenticChatInteractionType.StopChat, message)
422420
}
423421

@@ -475,7 +473,6 @@ export class ChatController {
475473

476474
private async processTabCloseMessage(message: TabClosedMessage) {
477475
this.sessionStorage.deleteSession(message.tabID)
478-
this.chatHistoryStorage.deleteHistory(message.tabID)
479476
this.triggerEventsStorage.removeTabEvents(message.tabID)
480477
// this.telemetryHelper.recordCloseChat(message.tabID)
481478
this.chatHistoryDb.updateTabOpenState(message.tabID, false)
@@ -1049,7 +1046,7 @@ export class ChatController {
10491046
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
10501047

10511048
this.sessionStorage.deleteSession(tabID)
1052-
this.chatHistoryStorage.getTabHistory(tabID).clear()
1049+
this.chatHistoryDb.clearTab(tabID)
10531050
}
10541051

10551052
private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1171,7 +1168,6 @@ export class ChatController {
11711168
switch (message.command) {
11721169
case 'clear':
11731170
this.sessionStorage.deleteSession(message.tabID)
1174-
this.chatHistoryStorage.getTabHistory(message.tabID).clear()
11751171
this.triggerEventsStorage.removeTabEvents(message.tabID)
11761172
recordTelemetryChatRunCommand('clear')
11771173
this.chatHistoryDb.clearTab(message.tabID)
@@ -1463,10 +1459,6 @@ export class ChatController {
14631459
}
14641460

14651461
const session = this.sessionStorage.getSession(tabID)
1466-
if (!session.localHistoryHydrated) {
1467-
triggerPayload.history = this.chatHistoryDb.getMessages(triggerEvent.tabID, 10)
1468-
session.localHistoryHydrated = true
1469-
}
14701462
await this.resolveContextCommandPayload(triggerPayload, session)
14711463
triggerPayload.useRelevantDocuments = triggerPayload.context.some(
14721464
(context) => typeof context !== 'string' && context.command === '@workspace'
@@ -1505,16 +1497,14 @@ export class ChatController {
15051497

15061498
const request = triggerPayloadToChatRequest(triggerPayload)
15071499

1508-
const chatHistory = this.chatHistoryStorage.getTabHistory(tabID)
15091500
const currentMessage = request.conversationState.currentMessage
15101501
if (currentMessage) {
1511-
chatHistory.fixHistory(currentMessage)
1502+
this.chatHistoryDb.fixHistory(tabID, currentMessage)
15121503
}
1513-
request.conversationState.history = chatHistory.getHistory()
1514-
1515-
const conversationId = chatHistory.getConversationId() || randomUUID()
1516-
chatHistory.setConversationId(conversationId)
1517-
request.conversationState.conversationId = conversationId
1504+
request.conversationState.history = this.chatHistoryDb
1505+
.getMessages(tabID)
1506+
.map((chat) => messageToChatMessage(chat))
1507+
request.conversationState.conversationId = session.sessionIdentifier
15181508

15191509
triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
15201510

@@ -1573,7 +1563,6 @@ export class ChatController {
15731563
this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
15741564

15751565
if (currentMessage && session.sessionIdentifier) {
1576-
chatHistory.appendUserMessage(currentMessage)
15771566
this.chatHistoryDb.addMessage(tabID, 'cwc', session.sessionIdentifier, {
15781567
body: triggerPayload.message,
15791568
type: 'prompt' as any,
@@ -1588,15 +1577,13 @@ export class ChatController {
15881577
response.$metadata.requestId
15891578
} metadata: ${inspect(response.$metadata, { depth: 12 })}`
15901579
)
1591-
15921580
this.cancelTokenSource = new vscode.CancellationTokenSource()
15931581
await this.messenger.sendAIResponse(
15941582
response,
15951583
session,
15961584
tabID,
15971585
triggerID,
15981586
triggerPayload,
1599-
chatHistory,
16001587
this.cancelTokenSource.token
16011588
)
16021589

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

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@ import {
6060
} from '@aws/mynah-ui'
6161
import { Database } from '../../../../shared/db/chatDb/chatDb'
6262
import { TabType } from '../../../../amazonq/webview/ui/storages/tabsStorage'
63-
import { ChatHistoryManager } from '../../../storages/chatHistory'
6463
import { ToolType, ToolUtils } from '../../../tools/toolUtils'
6564
import { ChatStream } from '../../../tools/chatStream'
6665
import path from 'path'
@@ -190,7 +189,6 @@ export class Messenger {
190189
tabID: string,
191190
triggerID: string,
192191
triggerPayload: TriggerPayload,
193-
chatHistoryManager: ChatHistoryManager,
194192
cancelToken: vscode.CancellationToken
195193
) {
196194
let message = ''
@@ -558,17 +556,6 @@ export class Messenger {
558556
)
559557
)
560558

561-
chatHistoryManager.pushAssistantMessage({
562-
assistantResponseMessage: {
563-
messageId: messageID,
564-
content: message,
565-
references: codeReference,
566-
...(toolUse &&
567-
toolUse.input !== undefined &&
568-
toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
569-
},
570-
})
571-
572559
getLogger().info(
573560
`All events received. requestId=%s counts=%s`,
574561
response.$metadata.requestId,

0 commit comments

Comments
 (0)