Skip to content

Commit 9df3b4c

Browse files
ctecte
committed
Tests + benchmarks for parseAssistantMessageV3
1 parent f18eb0e commit 9df3b4c

File tree

3 files changed

+167
-6
lines changed

3 files changed

+167
-6
lines changed

‎src/core/assistant-message/__tests__/parseAssistantMessage.test.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,12 @@ import { TextContent, ToolUse } from "../../../shared/tools"
44

55
import { AssistantMessageContent, parseAssistantMessage as parseAssistantMessageV1 } from "../parseAssistantMessage"
66
import { parseAssistantMessageV2 } from "../parseAssistantMessageV2"
7+
import { parseAssistantMessageV3 } from "../parseAssistantMessageV3"
78

89
const isEmptyTextContent = (block: AssistantMessageContent) =>
910
block.type === "text" && (block as TextContent).content === ""
1011

11-
;[parseAssistantMessageV1, parseAssistantMessageV2].forEach((parser, index) => {
12+
;[parseAssistantMessageV1, parseAssistantMessageV2, parseAssistantMessageV3].forEach((parser, index) => {
1213
describe(`parseAssistantMessageV${index + 1}`, () => {
1314
describe("text content parsing", () => {
1415
it("should parse a simple text message", () => {

‎src/core/assistant-message/__tests__/parseAssistantMessageBenchmark.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import { performance } from "perf_hooks"
44
import { parseAssistantMessage as parseAssistantMessageV1 } from "../parseAssistantMessage"
55
import { parseAssistantMessageV2 } from "../parseAssistantMessageV2"
6+
import { parseAssistantMessageV3 } from "../parseAssistantMessageV3"
67

78
const formatNumber = (num: number): string => {
89
return num.toString().replace(/\B(?=(\d{3})+(?!\d))/g, ",")
@@ -81,27 +82,29 @@ const runBenchmark = () => {
8182
const namePadding = maxNameLength + 2
8283

8384
console.log(
84-
`| ${"Test Case".padEnd(namePadding)} | V1 Time (ms) | V2 Time (ms) | V1/V2 Ratio | V1 Heap (bytes) | V2 Heap (bytes) |`,
85+
`| ${"Test Case".padEnd(namePadding)} | V1 Time (ms) | V2 Time (ms) | V3 Time (ms) | V1 Heap (bytes) | V2 Heap (bytes) | V3 Heap (bytes) |`,
8586
)
8687
console.log(
87-
`| ${"-".repeat(namePadding)} | ------------ | ------------ | ----------- | ---------------- | ---------------- |`,
88+
`| ${"-".repeat(namePadding)} | ------------ | ------------ | ------------ | ---------------- | ---------------- | ---------------- |`,
8889
)
8990

9091
for (const testCase of testCases) {
9192
const v1Time = measureExecutionTime(parseAssistantMessageV1, testCase.input)
9293
const v2Time = measureExecutionTime(parseAssistantMessageV2, testCase.input)
93-
const timeRatio = v1Time / v2Time
94+
const v3Time = measureExecutionTime(parseAssistantMessageV3, testCase.input)
9495

9596
const v1Memory = measureMemoryUsage(parseAssistantMessageV1, testCase.input)
9697
const v2Memory = measureMemoryUsage(parseAssistantMessageV2, testCase.input)
98+
const v3Memory = measureMemoryUsage(parseAssistantMessageV3, testCase.input)
9799

98100
console.log(
99101
`| ${testCase.name.padEnd(namePadding)} | ` +
100102
`${v1Time.toFixed(4).padStart(12)} | ` +
101103
`${v2Time.toFixed(4).padStart(12)} | ` +
102-
`${timeRatio.toFixed(2).padStart(11)} | ` +
104+
`${v3Time.toFixed(4).padStart(12)} | ` +
103105
`${formatNumber(Math.round(v1Memory.heapUsed)).padStart(16)} | ` +
104-
`${formatNumber(Math.round(v2Memory.heapUsed)).padStart(16)} |`,
106+
`${formatNumber(Math.round(v2Memory.heapUsed)).padStart(16)} | ` +
107+
`${formatNumber(Math.round(v3Memory.heapUsed)).padStart(16)} |`,
105108
)
106109
}
107110
}

‎src/core/assistant-message/parseAssistantMessageV3.ts‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
import { TextContent, ToolParamName, ToolUse, toolParamNames } from "../../shared/tools"
2+
import { toolNames, ToolName } from "../../schemas"
3+
4+
export type AssistantMessageContent = TextContent | ToolUse
5+
6+
// We keep parsing state between chunks so we don't repeatedly
7+
// re-scan the entire assistant message each time a new token arrives.
8+
export interface ParserState {
9+
currentToolUse?: ToolUse
10+
currentParamName?: ToolParamName
11+
currentTextContent?: TextContent
12+
}
13+
14+
const toolOpeningTags = toolNames.map((name) => `<${name}>`)
15+
const paramOpeningTags = toolParamNames.map((name) => `<${name}>`)
16+
17+
/**
18+
* Parses a chunk of streaming assistant text, mutating `state` and
19+
* returning *completed* content blocks in the order they were closed.
20+
*/
21+
export function parseAssistantMessageV3(chunk: string, state?: ParserState): AssistantMessageContent[] {
22+
if (!state) {
23+
state = createInitialParserState()
24+
}
25+
26+
let accumulator = ""
27+
const completedBlocks: AssistantMessageContent[] = []
28+
29+
// Finalise and push current text block.
30+
const flushText = () => {
31+
if (state.currentTextContent) {
32+
state.currentTextContent.content = state.currentTextContent.content + accumulator.trimEnd()
33+
accumulator = ""
34+
// Only push when a tool starts or chunk ends (done below).
35+
}
36+
}
37+
38+
for (let i = 0; i < chunk.length; i++) {
39+
const char = chunk[i]
40+
accumulator += char
41+
42+
if (state.currentToolUse && state.currentParamName) {
43+
// Inside a param – look for param closing tag.
44+
const paramClosing = `</${state.currentParamName}>`
45+
46+
if (accumulator.endsWith(paramClosing)) {
47+
// Strip closing tag from param value.
48+
const value = accumulator.slice(0, -paramClosing.length)
49+
50+
state.currentToolUse.params[state.currentParamName] =
51+
(state.currentToolUse.params[state.currentParamName] ?? "") + value
52+
53+
state.currentParamName = undefined
54+
accumulator = "" // Reset accumulator for next parsing portion.
55+
}
56+
57+
continue // Still inside param; don't treat anything else.
58+
}
59+
60+
// Inside a tool but not inside param..
61+
if (state.currentToolUse) {
62+
const toolClosing = `</${state.currentToolUse.name}>`
63+
64+
if (accumulator.endsWith(toolClosing)) {
65+
// Tool block complete.
66+
state.currentToolUse.partial = false
67+
state.currentToolUse = undefined
68+
accumulator = ""
69+
continue
70+
}
71+
72+
// Look for param opening tags.
73+
for (const open of paramOpeningTags) {
74+
if (accumulator.endsWith(open)) {
75+
state.currentParamName = open.slice(1, -1) as ToolParamName
76+
accumulator = "" // Reset; param value starts next char.
77+
// Initialize empty param value.
78+
state.currentToolUse.params[state.currentParamName] = ""
79+
break
80+
}
81+
}
82+
83+
continue
84+
}
85+
86+
// Not in a tool – look for opening of a tool.
87+
for (const open of toolOpeningTags) {
88+
if (accumulator.endsWith(open)) {
89+
// Flush any text accumulated *before* the tool tag.
90+
if (state.currentTextContent) {
91+
state.currentTextContent.partial = false
92+
completedBlocks.push(state.currentTextContent)
93+
state.currentTextContent = undefined
94+
}
95+
96+
state.currentToolUse = {
97+
type: "tool_use",
98+
name: open.slice(1, -1) as ToolName,
99+
params: {},
100+
partial: true,
101+
}
102+
103+
// Expose the new (partial) tool block immediately so the UI can
104+
// respond (e.g. opening diff view for write_to_file).
105+
completedBlocks.push(state.currentToolUse)
106+
accumulator = "" // Reset; tool content begins next char.
107+
break
108+
}
109+
}
110+
111+
// If we just started a tool, skip further text processing.
112+
if (state.currentToolUse) {
113+
continue
114+
}
115+
116+
// Regular text content.
117+
if (!state.currentTextContent) {
118+
state.currentTextContent = { type: "text", content: "", partial: true }
119+
120+
// Immediately push the partial text block so downstream consumers
121+
// can stream it progressively.
122+
completedBlocks.push(state.currentTextContent)
123+
}
124+
125+
// Text will be flushed at the end or before next block.
126+
}
127+
128+
// End of chunk – append remaining accumulator to current context.
129+
130+
if (state.currentParamName && state.currentToolUse) {
131+
// Still inside param value.
132+
state.currentToolUse.params[state.currentParamName] =
133+
(state.currentToolUse.params[state.currentParamName] ?? "") + accumulator
134+
135+
accumulator = ""
136+
} else if (state.currentToolUse) {
137+
// Inside tool but outside param – append to a special placeholder param? (ignored).
138+
// Nothing to do.
139+
} else {
140+
// Plain text.
141+
if (!state.currentTextContent) {
142+
state.currentTextContent = { type: "text", content: "", partial: true }
143+
}
144+
145+
state.currentTextContent.content += accumulator
146+
accumulator = ""
147+
}
148+
149+
// If end of chunk finalises a text block fully (no open tool starting later).
150+
flushText()
151+
152+
return completedBlocks
153+
}
154+
155+
export function createInitialParserState(): ParserState {
156+
return { currentToolUse: undefined, currentParamName: undefined, currentTextContent: undefined }
157+
}

0 commit comments

Comments
 (0)