From 4f51316f76dfcf5e72667dd9d12e53a91a430d02 Mon Sep 17 00:00:00 2001 From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com> Date: Sat, 5 Oct 2024 19:20:37 -0400 Subject: [PATCH] Refactor parseAssistantMessage --- src/core/ClaudeDev.ts | 147 ++------------------ src/core/prompts/parse-assistant-message.ts | 143 +++++++++++++++++++ 2 files changed, 151 insertions(+), 139 deletions(-) create mode 100644 src/core/prompts/parse-assistant-message.ts diff --git a/src/core/ClaudeDev.ts b/src/core/ClaudeDev.ts index 52dc40a..8e47726 100644 --- a/src/core/ClaudeDev.ts +++ b/src/core/ClaudeDev.ts @@ -49,6 +49,7 @@ import { formatResponse } from "./prompts/responses" import { addCustomInstructions, SYSTEM_PROMPT } from "./prompts/system" import { truncateHalfConversation } from "./sliding-window" import { ClaudeDevProvider, GlobalFileNames } from "./webview/ClaudeDevProvider" +import { parseAssistantMessage } from "./prompts/parse-assistant-message" const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0) ?? path.join(os.homedir(), "Desktop") // may or may not exist but fs checking existence would immediately ask for permission which would be bad UX, need to come up with a better solution @@ -1538,144 +1539,6 @@ export class ClaudeDev { } } - parseAssistantMessage(assistantMessage: string) { - let contentBlocks: AssistantMessageContent[] = [] - let currentTextContent: TextContent | undefined = undefined - let currentTextContentStartIndex = 0 - let currentToolUse: ToolUse | undefined = undefined - let currentToolUseStartIndex = 0 - let currentParamName: ToolParamName | undefined = undefined - let currentParamValueStartIndex = 0 - let accumulator = "" - - for (let i = 0; i < assistantMessage.length; i++) { - const char = assistantMessage[i] - accumulator += char - - // there should not be a param without a tool use - if (currentToolUse && currentParamName) { - const currentParamValue = accumulator.slice(currentParamValueStartIndex) - const paramClosingTag = `` - if (currentParamValue.endsWith(paramClosingTag)) { - // end of param value - currentToolUse.params[currentParamName] = currentParamValue.slice(0, -paramClosingTag.length).trim() - currentParamName = undefined - continue - } else { - // partial param value is accumulating - continue - } - } - - // no currentParamName - - if (currentToolUse) { - const currentToolValue = accumulator.slice(currentToolUseStartIndex) - const toolUseClosingTag = `` - if (currentToolValue.endsWith(toolUseClosingTag)) { - // end of a tool use - currentToolUse.partial = false - contentBlocks.push(currentToolUse) - currentToolUse = undefined - continue - } else { - const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`) - for (const paramOpeningTag of possibleParamOpeningTags) { - if (accumulator.endsWith(paramOpeningTag)) { - // start of a new parameter - currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName - currentParamValueStartIndex = accumulator.length - break - } - } - - // there's no current param, and not starting a new param - - // special case for write_to_file where file contents could contain the closing tag, in which case the param would have closed and we end up with the rest of the file contents here. To work around this, we get the string between the starting content tag and the LAST content tag. - const contentParamName: ToolParamName = "content" - if (currentToolUse.name === "write_to_file" && accumulator.endsWith(``)) { - const toolContent = accumulator.slice(currentToolUseStartIndex) - const contentStartTag = `<${contentParamName}>` - const contentEndTag = `` - const contentStartIndex = toolContent.indexOf(contentStartTag) + contentStartTag.length - const contentEndIndex = toolContent.lastIndexOf(contentEndTag) - if (contentStartIndex !== -1 && contentEndIndex !== -1 && contentEndIndex > contentStartIndex) { - currentToolUse.params[contentParamName] = toolContent - .slice(contentStartIndex, contentEndIndex) - .trim() - } - } - - // partial tool value is accumulating - continue - } - } - - // no currentToolUse - - let didStartToolUse = false - const possibleToolUseOpeningTags = toolUseNames.map((name) => `<${name}>`) - for (const toolUseOpeningTag of possibleToolUseOpeningTags) { - if (accumulator.endsWith(toolUseOpeningTag)) { - // start of a new tool use - currentToolUse = { - type: "tool_use", - name: toolUseOpeningTag.slice(1, -1) as ToolUseName, - params: {}, - partial: true, - } - currentToolUseStartIndex = accumulator.length - // this also indicates the end of the current text content - if (currentTextContent) { - currentTextContent.partial = false - // remove the partially accumulated tool use tag from the end of text ( prevLength) { - this.userMessageContentReady = false // new content we need to present, reset to false in case previous content set this to true - } - } - async recursivelyMakeClaudeRequests( userContent: UserContent, includeFileDetails: boolean = false @@ -1829,7 +1692,13 @@ export class ClaudeDev { break case "text": assistantMessage += chunk.text - this.parseAssistantMessage(assistantMessage) + // parse raw assistant message into content blocks + const prevLength = this.assistantMessageContent.length + this.assistantMessageContent = parseAssistantMessage(assistantMessage) + if (this.assistantMessageContent.length > prevLength) { + this.userMessageContentReady = false // new content we need to present, reset to false in case previous content set this to true + } + // present content to user this.presentAssistantMessage() break } diff --git a/src/core/prompts/parse-assistant-message.ts b/src/core/prompts/parse-assistant-message.ts new file mode 100644 index 0000000..84e7cac --- /dev/null +++ b/src/core/prompts/parse-assistant-message.ts @@ -0,0 +1,143 @@ +import { + AssistantMessageContent, + TextContent, + ToolUse, + ToolParamName, + toolParamNames, + toolUseNames, + ToolUseName, +} from "./AssistantMessage" + +export function parseAssistantMessage(assistantMessage: string) { + let contentBlocks: AssistantMessageContent[] = [] + let currentTextContent: TextContent | undefined = undefined + let currentTextContentStartIndex = 0 + let currentToolUse: ToolUse | undefined = undefined + let currentToolUseStartIndex = 0 + let currentParamName: ToolParamName | undefined = undefined + let currentParamValueStartIndex = 0 + let accumulator = "" + + for (let i = 0; i < assistantMessage.length; i++) { + const char = assistantMessage[i] + accumulator += char + + // there should not be a param without a tool use + if (currentToolUse && currentParamName) { + const currentParamValue = accumulator.slice(currentParamValueStartIndex) + const paramClosingTag = `` + if (currentParamValue.endsWith(paramClosingTag)) { + // end of param value + currentToolUse.params[currentParamName] = currentParamValue.slice(0, -paramClosingTag.length).trim() + currentParamName = undefined + continue + } else { + // partial param value is accumulating + continue + } + } + + // no currentParamName + + if (currentToolUse) { + const currentToolValue = accumulator.slice(currentToolUseStartIndex) + const toolUseClosingTag = `` + if (currentToolValue.endsWith(toolUseClosingTag)) { + // end of a tool use + currentToolUse.partial = false + contentBlocks.push(currentToolUse) + currentToolUse = undefined + continue + } else { + const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`) + for (const paramOpeningTag of possibleParamOpeningTags) { + if (accumulator.endsWith(paramOpeningTag)) { + // start of a new parameter + currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName + currentParamValueStartIndex = accumulator.length + break + } + } + + // there's no current param, and not starting a new param + + // special case for write_to_file where file contents could contain the closing tag, in which case the param would have closed and we end up with the rest of the file contents here. To work around this, we get the string between the starting content tag and the LAST content tag. + const contentParamName: ToolParamName = "content" + if (currentToolUse.name === "write_to_file" && accumulator.endsWith(``)) { + const toolContent = accumulator.slice(currentToolUseStartIndex) + const contentStartTag = `<${contentParamName}>` + const contentEndTag = `` + const contentStartIndex = toolContent.indexOf(contentStartTag) + contentStartTag.length + const contentEndIndex = toolContent.lastIndexOf(contentEndTag) + if (contentStartIndex !== -1 && contentEndIndex !== -1 && contentEndIndex > contentStartIndex) { + currentToolUse.params[contentParamName] = toolContent + .slice(contentStartIndex, contentEndIndex) + .trim() + } + } + + // partial tool value is accumulating + continue + } + } + + // no currentToolUse + + let didStartToolUse = false + const possibleToolUseOpeningTags = toolUseNames.map((name) => `<${name}>`) + for (const toolUseOpeningTag of possibleToolUseOpeningTags) { + if (accumulator.endsWith(toolUseOpeningTag)) { + // start of a new tool use + currentToolUse = { + type: "tool_use", + name: toolUseOpeningTag.slice(1, -1) as ToolUseName, + params: {}, + partial: true, + } + currentToolUseStartIndex = accumulator.length + // this also indicates the end of the current text content + if (currentTextContent) { + currentTextContent.partial = false + // remove the partially accumulated tool use tag from the end of text (