Optimize Parse_partial_json To Avoid O(N^2) Re-Parsing During Large Streaming Tool Calls
Streaming large tool call arguments through (prompt | model).stream() or with_structured_output(...).stream() triggers repeated parse_partial_json calls on every AIMessageChunk addition, causing quadratic CPU time as arguments grow.
AIMessageChunk.__add__ invokes parse_partial_json on the accumulated tool call args for each new chunk. The current parse_partial_json walks backward one character at a time and retries json.loads on each prefix, causing O(n^2) behavior for large argument strings.
Run the provided ToolCallStreamer with sizes 5000, 10000, 20000 and compare model.stream() vs (prompt | model).stream(). Observe model.stream() stays fast while the chained stream degrades superlinearly (e.g. 0.11s vs 2.06s for 20k chars).
Fixing Code Block
import json
import re
from typing import Any
_STRING_RE = re.compile(r'"(?:[^"\\]|\\.)*"')
def parse_partial_json(s: str, *, strict: bool = False) -> Any:
new_s = s.strip()
if new_s.startswith('[') and not new_s.endswith(']'):
new_s += ']'
if new_s.startswith('{') and not new_s.endswith('}'):
new_s += '}'
try:
return json.loads(new_s, strict=strict)
except json.JSONDecodeError as e:
error_pos = e.pos - 1
for match in _STRING_RE.finditer(new_s):
if match.start() <= error_pos < match.end():
error_pos = match.start()
break
candidate = new_s[:error_pos].rstrip()
try:
return json.loads(candidate, strict=strict)
except json.JSONDecodeError:
pass
for i in range(error_pos, 0, -1):
candidate = new_s[:i].rstrip()
try:
return json.loads(candidate, strict=strict)
except json.JSONDecodeError:
continue
raise
The hotfix uses json.JSONDecodeError.pos to jump to the first invalid character instead of trimming from the end. It also scans string literals with a compiled regex so an unterminated string is cut at its opening quote, preserving the longest valid prefix. The fallback loop only examines the small prefix before the error position, avoiding O(n^2) retries.
Edge Case Audit
This changes the internal trimming strategy; in rare malformed JSON cases the fallback may still degrade or return a slightly different partial object than the legacy loop. Add unit tests for partial objects, arrays, escaped quotes, and unterminated strings before deploying. If regressions appear, revert to the previous character-trim implementation. The regex is module-level and immutable, so no threading concerns, but performance for deeply nested or extremely large JSON may still require further chunk-merging optimizations.
