Refactoring the fix history fucntion

Author: ashishrp-aws

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

@@ -7,6 +7,7 @@ import {
     Tool,
     ToolResult,
     ToolResultStatus,
+    ToolUse,
     UserInputMessage,
     UserInputMessageContext,
 } from '@amzn/codewhisperer-streaming'
@@ -105,115 +106,99 @@ export class ChatHistoryManager {
      *    message is set without tool results, then the user message will have cancelled tool results.
      */
     public fixHistory(newUserMessage: ChatMessage): ChatMessage {
-        // Trim the conversation history if it exceeds the maximum length
-        if (this.history.length > MaxConversationHistoryLength) {
-            // Find the second oldest user message without tool results
-            let indexToTrim: number | undefined
-
-            for (let i = 1; i < this.history.length; i++) {
-                const message = this.history[i]
-                if (message.userInputMessage) {
-                    const userMessage = message.userInputMessage
-                    const ctx = userMessage.userInputMessageContext
-                    const hasNoToolResults = ctx && (!ctx.toolResults || ctx.toolResults.length === 0)
-                    if (hasNoToolResults && userMessage.content !== '') {
-                        indexToTrim = i
-                        break
-                    }
-                }
-            }
-            if (indexToTrim !== undefined) {
-                this.logger.debug(`Removing the first ${indexToTrim} elements in the history`)
-                this.history.splice(0, indexToTrim)
-            } else {
-                this.logger.debug('No valid starting user message found in the history, clearing')
-                this.history = []
+        this.trimConversationHistory()
+        this.ensureLastMessageFromAssistant()
+        return this.handleToolUses(newUserMessage)
+    }
+
+    private trimConversationHistory(): void {
+        if (this.history.length <= MaxConversationHistoryLength) {
+            return
+        }
+
+        const indexToTrim = this.findIndexToTrim()
+        if (indexToTrim !== undefined) {
+            this.logger.debug(`Removing the first ${indexToTrim} elements in the history`)
+            this.history.splice(0, indexToTrim)
+        } else {
+            this.logger.debug('No valid starting user message found in the history, clearing')
+            this.history = []
+        }
+    }
+
+    private findIndexToTrim(): number | undefined {
+        for (let i = 1; i < this.history.length; i++) {
+            const message = this.history[i]
+            if (this.isValidUserMessageWithoutToolResults(message)) {
+                return i
             }
         }
+        return undefined
+    }
 
-        // Ensure the last message is from the assistant
+    private isValidUserMessageWithoutToolResults(message: ChatMessage): boolean {
+        if (!message.userInputMessage) {
+            return false
+        }
+        const ctx = message.userInputMessage.userInputMessageContext
+        return Boolean(
+            ctx && (!ctx.toolResults || ctx.toolResults.length === 0) && message.userInputMessage.content !== ''
+        )
+    }
+
+    private ensureLastMessageFromAssistant(): void {
         if (this.history.length > 0 && this.history[this.history.length - 1].userInputMessage !== undefined) {
             this.logger.debug('Last message in history is from the user, dropping')
             this.history.pop()
         }
+    }
 
-        // If the last message from the assistant contains tool uses, ensure the next user message contains tool results
-
+    private handleToolUses(newUserMessage: ChatMessage): ChatMessage {
         const lastHistoryMessage = this.history[this.history.length - 1]
+        if (!lastHistoryMessage || !lastHistoryMessage.assistantResponseMessage || !newUserMessage) {
+            return newUserMessage
+        }
 
-        if (
-            lastHistoryMessage &&
-            (lastHistoryMessage.assistantResponseMessage ||
-                lastHistoryMessage.assistantResponseMessage !== undefined) &&
-            newUserMessage
-        ) {
-            const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
-
-            if (toolUses && toolUses.length > 0) {
-                if (newUserMessage.userInputMessage) {
-                    if (newUserMessage.userInputMessage.userInputMessageContext) {
-                        const ctx = newUserMessage.userInputMessage.userInputMessageContext
-
-                        if (!ctx.toolResults || ctx.toolResults.length === 0) {
-                            ctx.toolResults = toolUses.map((toolUse) => ({
-                                toolUseId: toolUse.toolUseId,
-                                content: [
-                                    {
-                                        type: 'Text',
-                                        text: 'Tool use was cancelled by the user',
-                                    },
-                                ],
-                                status: ToolResultStatus.ERROR,
-                            }))
-                        }
-                    } else {
-                        const toolResults = toolUses.map((toolUse) => ({
-                            toolUseId: toolUse.toolUseId,
-                            content: [
-                                {
-                                    type: 'Text',
-                                    text: 'Tool use was cancelled by the user',
-                                },
-                            ],
-                            status: ToolResultStatus.ERROR,
-                        }))
-
-                        newUserMessage.userInputMessage.userInputMessageContext = {
-                            shellState: undefined,
-                            envState: undefined,
-                            toolResults: toolResults,
-                            tools: this.tools.length === 0 ? undefined : [...this.tools],
-                        }
-
-                        return newUserMessage
-                    }
-                }
-            }
+        const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
+        if (!toolUses || toolUses.length === 0) {
+            return newUserMessage
         }
 
-        // Always return the message to fix the TypeScript error
-        return newUserMessage
+        return this.addToolResultsToUserMessage(newUserMessage, toolUses)
     }
 
-    /**
-     * Adds tool results to the conversation.
-     */
-    addToolResults(toolResults: ToolResult[]): void {
-        const userInputMessageContext: UserInputMessageContext = {
-            shellState: undefined,
-            envState: undefined,
-            toolResults: toolResults,
-            tools: this.tools.length === 0 ? undefined : [...this.tools],
+    private addToolResultsToUserMessage(newUserMessage: ChatMessage, toolUses: ToolUse[]): ChatMessage {
+        if (!newUserMessage.userInputMessage) {
+            return newUserMessage
         }
 
-        const msg: UserInputMessage = {
-            content: '',
-            userInputMessageContext: userInputMessageContext,
-        }
+        const toolResults = this.createToolResults(toolUses)
 
-        if (this.lastUserMessage?.userInputMessage) {
-            this.lastUserMessage.userInputMessage = msg
+        if (newUserMessage.userInputMessage.userInputMessageContext) {
+            newUserMessage.userInputMessage.userInputMessageContext.toolResults = toolResults
+        } else {
+            newUserMessage.userInputMessage.userInputMessageContext = {
+                shellState: undefined,
+                envState: undefined,
+                toolResults: toolResults,
+                tools: this.tools.length === 0 ? undefined : [...this.tools],
+            }
         }
+
+        return newUserMessage
+    }
+
+    private createToolResults(toolUses: ToolUse[]): ToolResult[] {
+        return toolUses.map((toolUse) => ({
+            toolUseId: toolUse.toolUseId,
+            content: [
+                {
+                    type: 'Text',
+                    text: 'Tool use was cancelled by the user',
+                },
+            ],
+            status: ToolResultStatus.ERROR,
+        }))
     }
 
     private formatChatHistoryMessage(message: ChatMessage): ChatMessage {