Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions packages/core/src/codewhispererChat/controllers/chat/controller.ts
Original file line number Diff line number Diff line change
Expand Up @@ -744,6 +744,7 @@ export class ChatController {
const toolResults: ToolResult[] = []

let response = ''
let shouldDisplayMessage = true
if (toolUseError) {
toolResults.push({
content: [{ text: toolUseError.message }],
Expand All @@ -753,6 +754,7 @@ export class ChatController {
if (toolUseError instanceof SyntaxError) {
response =
"Your toolUse input isn't valid. Please check the syntax and make sure the input is complete. If the input is large, break it down into multiple tool uses with smaller input."
shouldDisplayMessage = false
}
} else {
const result = ToolUtils.tryFromToolUse(toolUse)
Expand Down Expand Up @@ -844,6 +846,7 @@ export class ChatController {
contextLengths: {
...defaultContextLengths,
},
shouldDisplayMessage: shouldDisplayMessage,
},
triggerID
)
Expand Down Expand Up @@ -1138,9 +1141,6 @@ export class ChatController {

this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)

this.sessionStorage.deleteSession(tabID)
this.chatHistoryDb.clearTab(tabID)
}

private async processContextMenuCommand(command: EditorContextCommand) {
Expand Down Expand Up @@ -1317,6 +1317,7 @@ export class ChatController {
this.processException(e, message.tabID)
}
}

private initialCleanUp(session: ChatSession) {
// Create a fresh token for this new conversation
session.createNewTokenSource()
Expand Down Expand Up @@ -1595,11 +1596,12 @@ export class ChatController {

const currentMessage = request.conversationState.currentMessage
if (currentMessage) {
this.chatHistoryDb.fixHistory(tabID, currentMessage)
this.chatHistoryDb.fixHistory(tabID, currentMessage, session.sessionIdentifier)
}
request.conversationState.history = this.chatHistoryDb
.getMessages(tabID)
.map((chat) => messageToChatMessage(chat))
// Do not include chatHistory for requests going to Mynah
request.conversationState.history = request.conversationState.currentMessage?.userInputMessage?.userIntent
? []
: this.chatHistoryDb.getMessages(tabID).map((chat) => messageToChatMessage(chat))
request.conversationState.conversationId = session.sessionIdentifier

triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
Expand Down Expand Up @@ -1669,6 +1671,7 @@ export class ChatController {
userIntent: currentMessage.userInputMessage?.userIntent,
origin: currentMessage.userInputMessage?.origin,
userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext,
shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
})
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,7 @@ export class Messenger {
toolUse && toolUse.input !== undefined && toolUse.input !== ''
? [{ ...toolUse }]
: undefined,
shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
})
}
if (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,7 @@ export interface TriggerPayload {
origin?: Origin
pairProgrammingModeOn?: boolean
history?: Message[]
shouldDisplayMessage?: boolean
}

export type ContextLengths = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,9 @@ export class TabBarController {
selectedTab.historyId,
selectedTab.tabType,
selectedTab.conversations.flatMap((conv: Conversation) =>
conv.messages.map((message) => messageToChatItem(message))
conv.messages
.filter((message) => message.shouldDisplayMessage !== false)
.map((message) => messageToChatItem(message))
),
exportTab
)
Expand Down
83 changes: 35 additions & 48 deletions packages/core/src/shared/db/chatDb/chatDb.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming'

// Maximum number of characters to keep in history
const MaxConversationHistoryCharacters = 600_000
// Maximum number of messages to keep in history
const MaxConversationHistoryMessages = 250

/**
* A singleton database class that manages chat history persistence using LokiJS.
Expand Down Expand Up @@ -144,18 +146,6 @@ export class Database {
return undefined
}

getActiveConversation(historyId: string): Conversation | undefined {
const tabCollection = this.db.getCollection<Tab>(TabCollection)
const tabData = tabCollection.findOne({ historyId })

if (!tabData?.conversations.length) {
this.logger.debug('No active conversations found')
return undefined
}

return tabData.conversations[0]
}

clearTab(tabId: string) {
if (this.initialized) {
const tabCollection = this.db.getCollection<Tab>(TabCollection)
Expand All @@ -170,34 +160,6 @@ export class Database {
}
}

// Removes the most recent message(s) from the chat history for a given tab
clearRecentHistory(tabId: string): void {
if (this.initialized) {
const historyId = this.historyIdMapping.get(tabId)
this.logger.info(`Clearing recent history: tabId=${tabId}, historyId=${historyId || 'undefined'}`)
if (historyId) {
const tabCollection = this.db.getCollection<Tab>(TabCollection)
const tabData = tabCollection.findOne({ historyId })
if (tabData) {
const activeConversation = tabData.conversations[0]
const allMessages = this.getMessages(tabId)
const lastMessage = allMessages[allMessages.length - 1]
this.logger.debug(`Last message type: ${lastMessage.type}`)

if (lastMessage.type === ('prompt' as ChatItemType)) {
allMessages.pop()
this.logger.info(`Removed last user message`)
} else {
allMessages.splice(-2)
this.logger.info(`Removed last assistant message and user message`)
}
activeConversation.messages = allMessages
tabCollection.update(tabData)
}
}
}
}

updateTabOpenState(tabId: string, isOpen: boolean) {
if (this.initialized) {
const tabCollection = this.db.getCollection<Tab>(TabCollection)
Expand Down Expand Up @@ -348,13 +310,14 @@ export class Database {

/**
* Fixes the history to maintain the following invariants:
* 1. The history character length is <= MAX_CONVERSATION_HISTORY_CHARACTERS. Oldest messages are dropped.
* 2. The first message is from the user. Oldest messages are dropped if needed.
* 3. The last message is from the assistant. The last message is dropped if it is from the user.
* 4. If the last message is from the assistant and it contains tool uses, and a next user
* 1. The history contains at most MaxConversationHistoryMessages messages. Oldest messages are dropped.
* 2. The history character length is <= MaxConversationHistoryCharacters. Oldest messages are dropped.
* 3. The first message is from the user. Oldest messages are dropped if needed.
* 4. The last message is from the assistant. The last message is dropped if it is from the user.
* 5. If the last message is from the assistant and it contains tool uses, and a next user
* message is set without tool results, then the user message will have cancelled tool results.
*/
fixHistory(tabId: string, newUserMessage: ChatMessage): void {
fixHistory(tabId: string, newUserMessage: ChatMessage, conversationId: string): void {
if (!this.initialized) {
return
}
Expand All @@ -371,10 +334,12 @@ export class Database {
return
}

const activeConversation = tabData.conversations[0]
let allMessages = activeConversation.messages
let allMessages = tabData.conversations.flatMap((conversation: Conversation) => conversation.messages)
this.logger.info(`Found ${allMessages.length} messages in conversation`)

// Make sure we don't exceed MaxConversationHistoryMessages
allMessages = this.trimHistoryToMaxLength(allMessages)

// Drop empty assistant partial if it’s the last message
this.handleEmptyAssistantMessage(allMessages)

Expand All @@ -387,11 +352,33 @@ export class Database {
// If the last message is from the assistant and it contains tool uses, and a next user message is set without tool results, then the user message will have cancelled tool results.
this.handleToolUses(allMessages, newUserMessage)

activeConversation.messages = allMessages
tabData.conversations = [
{
conversationId: conversationId,
clientType: ClientType.VSCode,
messages: allMessages,
},
]
tabData.updatedAt = new Date()
tabCollection.update(tabData)
this.logger.info(`Updated tab data in collection`)
}

private trimHistoryToMaxLength(messages: Message[]): Message[] {
while (messages.length > MaxConversationHistoryMessages) {
// Find the next valid user message to start from
const indexToTrim = this.findIndexToTrim(messages)
if (indexToTrim !== undefined && indexToTrim > 0) {
this.logger.debug(`Removing the first ${indexToTrim} elements to maintain valid history length`)
messages.splice(0, indexToTrim)
} else {
this.logger.debug('Could not find a valid point to trim, reset history to reduce history size')
return []
}
}
return messages
}

private handleEmptyAssistantMessage(messages: Message[]): void {
if (messages.length === 0) {
return
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/shared/db/chatDb/util.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ export type Message = {
origin?: Origin
userInputMessageContext?: UserInputMessageContext
toolUses?: ToolUse[]
shouldDisplayMessage?: boolean
}

/**
Expand Down
Loading