From 51f236dce2272e1fed2aa228f13ae18dbf463990 Mon Sep 17 00:00:00 2001 From: Afonso Ribeiro Date: Fri, 16 Oct 2026 18:42:10 +0100 Subject: [PATCH] perf(streaming): skip re-parsing tool input JSON while a delta only extends an open string --- src/anthropic/lib/streaming/_beta_messages.py | 21 ++-- src/anthropic/lib/streaming/_messages.py | 19 ++-- src/anthropic/lib/streaming/_partial_json.py | 97 +++++++++++++++++++ tests/lib/streaming/helpers.py | 82 +++++++++++++++- tests/lib/streaming/test_beta_messages.py | 43 +++++++- tests/lib/streaming/test_messages.py | 58 ++++++++++- 6 files changed, 296 insertions(+), 24 deletions(-) create mode 100644 src/anthropic/lib/streaming/_partial_json.py diff --git a/src/anthropic/lib/streaming/_beta_messages.py b/src/anthropic/lib/streaming/_beta_messages.py index 07ae383a5..5b8afab72 100644 --- a/src/anthropic/lib/streaming/_beta_messages.py +++ b/src/anthropic/lib/streaming/_beta_messages.py @@ -29,6 +29,7 @@ ) from ..._streaming import Stream, AsyncStream from ...types.beta import BetaRawMessageStreamEvent +from ._partial_json import PartialJsonBuffer from ..._utils._utils import is_given from .._parse._response import ResponseFormatT, parse_text from ...types.beta.parsed_beta_message import ParsedBetaMessage, ParsedBetaContentBlock @@ -54,7 +55,7 @@ def __init__( self.text_stream = self.__stream_text__() self._iterator = self.__stream__() self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None - self.__json_bufs: dict[int, bytes] = {} + self.__json_bufs: dict[int, PartialJsonBuffer] = {} self.__output_format = output_format share_message_snapshot(raw_stream, self._current_snapshot_or_none) @@ -220,7 +221,7 @@ def __init__( self.text_stream = self.__stream_text__() self._iterator = self.__stream__() self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None - self.__json_bufs: dict[int, bytes] = {} + self.__json_bufs: dict[int, PartialJsonBuffer] = {} self.__output_format = output_format share_message_snapshot(raw_stream, self._current_snapshot_or_none) @@ -481,7 +482,7 @@ def accumulate_event( *, event: BetaRawMessageStreamEvent, current_snapshot: ParsedBetaMessage[ResponseFormatT] | None, - json_bufs: dict[int, bytes], + json_bufs: dict[int, PartialJsonBuffer], request_headers: httpx2.Headers, output_format: ResponseFormatT | NotGiven = not_given, ) -> ParsedBetaMessage[ResponseFormatT]: @@ -527,23 +528,21 @@ def accumulate_event( if isinstance(content, TRACKS_TOOL_INPUT): from jiter import from_json - json_buf = json_bufs.get(event.index, b"") - json_buf += bytes(event.delta.partial_json, "utf-8") + json_buf = json_bufs.setdefault(event.index, PartialJsonBuffer()) + extends_open_string = json_buf.feed(bytes(event.delta.partial_json, "utf-8")) if json_buf: try: anthropic_beta = request_headers.get("anthropic-beta", "") if request_headers else "" if "fine-grained-tool-streaming-2025-05-14" in anthropic_beta: - content.input = from_json(json_buf, partial_mode="trailing-strings") - else: - content.input = from_json(json_buf, partial_mode=True) + content.input = from_json(json_buf.getvalue(), partial_mode="trailing-strings") + elif not extends_open_string: + content.input = from_json(json_buf.getvalue(), partial_mode=True) except ValueError as e: raise ValueError( - f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.decode('utf-8')}" + f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.getvalue().decode('utf-8')}" ) from e - - json_bufs[event.index] = json_buf elif event.delta.type == "citations_delta": if content.type == "text": if not content.citations: diff --git a/src/anthropic/lib/streaming/_messages.py b/src/anthropic/lib/streaming/_messages.py index 8ef8721d3..e12912224 100644 --- a/src/anthropic/lib/streaming/_messages.py +++ b/src/anthropic/lib/streaming/_messages.py @@ -26,6 +26,7 @@ from ..._models import build, construct_type, construct_type_unchecked from .._tracing import share_message_snapshot, withhold_message_snapshot from ..._streaming import Stream, AsyncStream +from ._partial_json import PartialJsonBuffer from ..._utils._utils import is_given from .._parse._response import ResponseFormatT, parse_text from ...types.parsed_message import ParsedMessage, ParsedContentBlock @@ -51,7 +52,7 @@ def __init__( self.text_stream = self.__stream_text__() self._iterator = self.__stream__() self.__final_message_snapshot: ParsedMessage[ResponseFormatT] | None = None - self.__json_bufs: dict[int, bytes] = {} + self.__json_bufs: dict[int, PartialJsonBuffer] = {} self.__output_format = output_format share_message_snapshot(raw_stream, self._current_snapshot_or_none) @@ -216,7 +217,7 @@ def __init__( self.text_stream = self.__stream_text__() self._iterator = self.__stream__() self.__final_message_snapshot: ParsedMessage[ResponseFormatT] | None = None - self.__json_bufs: dict[int, bytes] = {} + self.__json_bufs: dict[int, PartialJsonBuffer] = {} self.__output_format = output_format share_message_snapshot(raw_stream, self._current_snapshot_or_none) @@ -465,7 +466,7 @@ def accumulate_event( *, event: RawMessageStreamEvent, current_snapshot: ParsedMessage[ResponseFormatT] | None, - json_bufs: dict[int, bytes], + json_bufs: dict[int, PartialJsonBuffer], output_format: ResponseFormatT | NotGiven = not_given, ) -> ParsedMessage[ResponseFormatT]: if not isinstance(cast(Any, event), BaseModel): @@ -502,18 +503,16 @@ def accumulate_event( if isinstance(content, TRACKS_TOOL_INPUT): from jiter import from_json - json_buf = json_bufs.get(event.index, b"") - json_buf += bytes(event.delta.partial_json, "utf-8") + json_buf = json_bufs.setdefault(event.index, PartialJsonBuffer()) + extends_open_string = json_buf.feed(bytes(event.delta.partial_json, "utf-8")) - if json_buf: + if json_buf and not extends_open_string: try: - content.input = from_json(json_buf, partial_mode=True) + content.input = from_json(json_buf.getvalue(), partial_mode=True) except ValueError as e: raise ValueError( - f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.decode('utf-8')}" + f"Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: {e}. JSON: {json_buf.getvalue().decode('utf-8')}" ) from e - - json_bufs[event.index] = json_buf elif event.delta.type == "citations_delta": if content.type == "text": if not content.citations: diff --git a/src/anthropic/lib/streaming/_partial_json.py b/src/anthropic/lib/streaming/_partial_json.py new file mode 100644 index 000000000..6cdf63085 --- /dev/null +++ b/src/anthropic/lib/streaming/_partial_json.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import re + +# Bytes that end a JSON string, start an escape or are not allowed in one. +_STRING_SPECIAL = re.compile(rb'["\\\x00-\x1f]') +_SIMPLE_ESCAPES = frozenset(b'"\\/bfnrt') +_HEX_DIGITS = frozenset(b"0123456789abcdefABCDEF") + + +class PartialJsonBuffer: + """Collects the `partial_json` deltas of one content block. + + `from_json(..., partial_mode=True)` drops a trailing unclosed string, so a delta that only extends + an open string does not change the parsed value and the buffer does not need to be parsed again. + Anything that is not valid inside a JSON string marks the buffer as unsafe to skip from then on, + so that the parser still sees it and raises the same error. + """ + + def __init__(self) -> None: + self._buf = bytearray() + self._in_string = False + self._invalid = False + # 0: no escape, 1: after a backslash, 2-5: reading the hex digits of `\u` + self._escape = 0 + self._code = 0 + self._high_surrogate = False + + def __len__(self) -> int: + return len(self._buf) + + def getvalue(self) -> bytes: + return bytes(self._buf) + + def feed(self, data: bytes) -> bool: + """Appends `data` and returns whether it only extended an open string with valid content.""" + self._buf += data + if self._invalid: + return False + + extends_open_string = self._in_string + pos = 0 + end = len(data) + while pos < end: + if not self._in_string: + quote = data.find(b'"', pos) + if quote < 0: + break + self._in_string = True + pos = quote + 1 + continue + + char = data[pos] + if self._escape == 0: + if self._high_surrogate and char != 0x5C: # a high surrogate must be followed by `\u` + self._invalid = True + return False + match = _STRING_SPECIAL.search(data, pos) + if match is None: + break + pos = match.end() + char = data[pos - 1] + if char == 0x22: # '"' + self._in_string = False + extends_open_string = False + elif char == 0x5C: # '\\' + self._escape = 1 + else: + self._invalid = True + return False + elif self._escape == 1: + pos += 1 + if char == 0x75: # 'u' + self._escape = 2 + self._code = 0 + elif char in _SIMPLE_ESCAPES and not self._high_surrogate: + self._escape = 0 + else: + self._invalid = True + return False + else: + pos += 1 + if char not in _HEX_DIGITS: + self._invalid = True + return False + self._code = self._code * 16 + int(chr(char), 16) + self._escape += 1 + if self._escape == 6: + self._escape = 0 + is_high = 0xD800 <= self._code <= 0xDBFF + is_low = 0xDC00 <= self._code <= 0xDFFF + if self._high_surrogate != is_low: + self._invalid = True + return False + self._high_surrogate = is_high + + return extends_open_string diff --git a/tests/lib/streaming/helpers.py b/tests/lib/streaming/helpers.py index bddcadf6e..aff70845e 100644 --- a/tests/lib/streaming/helpers.py +++ b/tests/lib/streaming/helpers.py @@ -1,9 +1,13 @@ from __future__ import annotations import os -from typing import TypeVar, Iterator +import json +from typing import Any, TypeVar, Iterator from typing_extensions import AsyncIterator +import jiter +import pytest + _T = TypeVar("_T") @@ -26,3 +30,79 @@ async def to_async_iter(iter: Iterator[_T]) -> AsyncIterator[_T]: """Convert a synchronous iterator to an asynchronous one.""" for event in iter: yield event + + +LONG_TOOL_INPUT = json.dumps({"path": "notes.md", "content": 'a "quoted" café line\\ \n' * 2000}) + + +def get_long_tool_input_response(tool_input: str = LONG_TOOL_INPUT) -> Iterator[bytes]: + """A `tool_use` block whose input holds a long string, streamed in small `input_json_delta` events.""" + events: list[dict[str, Any]] = [ + { + "type": "message_start", + "message": { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_01", + "name": "write_file", + "input": {}, + "caller": {"type": "direct"}, + }, + }, + *( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": tool_input[i : i + 8]}, + } + for i in range(0, len(tool_input), 8) + ), + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 100}, + }, + {"type": "message_stop"}, + ] + for event in events: + yield f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def record_parsed_json_sizes(monkeypatch: pytest.MonkeyPatch) -> list[int]: + """Patches `jiter.from_json` to record the size of every input it parses.""" + sizes: list[int] = [] + from_json = jiter.from_json + + def recording_from_json(json_data: bytes, **kwargs: Any) -> Any: + sizes.append(len(json_data)) + return from_json(json_data, **kwargs) + + monkeypatch.setattr(jiter, "from_json", recording_from_json) + return sizes + + +def assert_long_tool_input_parsed(snapshots: list[object], final_input: object, parsed_sizes: list[int]) -> None: + # re-parsing the whole buffer on every delta would add up to hundreds of times the input size + assert sum(parsed_sizes) < 10 * len(LONG_TOOL_INPUT) + + expected: list[object] = [] + for i in range(8, len(LONG_TOOL_INPUT), 8): + expected.append(jiter.from_json(LONG_TOOL_INPUT[:i].encode(), partial_mode=True)) + expected.append(json.loads(LONG_TOOL_INPUT)) + + assert snapshots == expected + assert final_input == json.loads(LONG_TOOL_INPUT) diff --git a/tests/lib/streaming/test_beta_messages.py b/tests/lib/streaming/test_beta_messages.py index 982d40f70..eead6167d 100644 --- a/tests/lib/streaming/test_beta_messages.py +++ b/tests/lib/streaming/test_beta_messages.py @@ -26,7 +26,13 @@ from anthropic.types.beta.beta_message_delta_usage import BetaMessageDeltaUsage from anthropic.types.beta.beta_raw_message_delta_event import Delta as BetaRawMessageDelta, BetaRawMessageDeltaEvent -from .helpers import get_response, to_async_iter +from .helpers import ( + get_response, + to_async_iter, + record_parsed_json_sizes, + get_long_tool_input_response, + assert_long_tool_input_parsed, +) base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") api_key = "my-anthropic-api-key" @@ -474,6 +480,22 @@ def test_tool_use(self, respx_mock: MockRouter) -> None: assert json.loads(respx_mock.calls.last.request.content)["tools"] == [WeatherTool().to_dict()] + @pytest.mark.respx(base_url=base_url) + def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + respx_mock.post("/v1/messages").mock(return_value=httpx2.Response(200, content=get_long_tool_input_response())) + parsed_sizes = record_parsed_json_sizes(monkeypatch) + + with sync_client.beta.messages.stream( + max_tokens=1024, + messages=[{"role": "user", "content": "Write the notes file"}], + model="claude-sonnet-5-5", + ) as stream: + snapshots = [event.snapshot for event in stream if event.type == "input_json"] + tool_use = stream.get_final_message().content[0] + + assert tool_use.type == "tool_use" + assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes) + @pytest.mark.respx(base_url=base_url) def test_tool_use_round_trip(self, respx_mock: MockRouter) -> None: route = respx_mock.post("/v1/messages").mock( @@ -838,6 +860,25 @@ async def test_tool_use(self, respx_mock: MockRouter) -> None: assert json.loads(respx_mock.calls.last.request.content)["tools"] == [WeatherTool().to_dict()] + @pytest.mark.asyncio + @pytest.mark.respx(base_url=base_url) + async def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + respx_mock.post("/v1/messages").mock( + return_value=httpx2.Response(200, content=to_async_iter(get_long_tool_input_response())) + ) + parsed_sizes = record_parsed_json_sizes(monkeypatch) + + async with async_client.beta.messages.stream( + max_tokens=1024, + messages=[{"role": "user", "content": "Write the notes file"}], + model="claude-sonnet-5-5", + ) as stream: + snapshots = [event.snapshot async for event in stream if event.type == "input_json"] + tool_use = (await stream.get_final_message()).content[0] + + assert tool_use.type == "tool_use" + assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes) + @pytest.mark.asyncio @pytest.mark.respx(base_url=base_url) async def test_server_tool_use(self, respx_mock: MockRouter) -> None: diff --git a/tests/lib/streaming/test_messages.py b/tests/lib/streaming/test_messages.py index 90d95ee37..06925d74d 100644 --- a/tests/lib/streaming/test_messages.py +++ b/tests/lib/streaming/test_messages.py @@ -21,7 +21,13 @@ from anthropic.types.message_delta_usage import MessageDeltaUsage from anthropic.types.raw_message_delta_event import Delta as RawMessageDelta, RawMessageDeltaEvent -from .helpers import get_response, to_async_iter +from .helpers import ( + get_response, + to_async_iter, + record_parsed_json_sizes, + get_long_tool_input_response, + assert_long_tool_input_parsed, +) base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") api_key = "my-anthropic-api-key" @@ -383,6 +389,37 @@ def test_tool_use_invalid_json(self, respx_mock: MockRouter) -> None: ) as stream: stream.until_done() + @pytest.mark.respx(base_url=base_url) + def test_tool_use_invalid_escape_in_long_string(self, respx_mock: MockRouter) -> None: + tool_input = '{"content": "' + "line " * 20 + "\\q" + "line " * 20 + '"}' + respx_mock.post("/v1/messages").mock( + return_value=httpx2.Response(200, content=get_long_tool_input_response(tool_input)) + ) + + with pytest.raises(ValueError, match=r"^Unable to parse tool parameter JSON from model.*invalid escape"): + with sync_client.messages.stream( + max_tokens=1024, + messages=[{"role": "user", "content": "Write the notes file"}], + model="claude-sonnet-5-5", + ) as stream: + stream.until_done() + + @pytest.mark.respx(base_url=base_url) + def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + respx_mock.post("/v1/messages").mock(return_value=httpx2.Response(200, content=get_long_tool_input_response())) + parsed_sizes = record_parsed_json_sizes(monkeypatch) + + with sync_client.messages.stream( + max_tokens=1024, + messages=[{"role": "user", "content": "Write the notes file"}], + model="claude-sonnet-5-5", + ) as stream: + snapshots = [event.snapshot for event in stream if event.type == "input_json"] + tool_use = stream.get_final_message().content[0] + + assert tool_use.type == "tool_use" + assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes) + @pytest.mark.respx(base_url=base_url) def test_tool_use_caller_omitted(self, respx_mock: MockRouter) -> None: respx_mock.post("/v1/messages").mock( @@ -690,6 +727,25 @@ async def test_tool_use_invalid_json(self, respx_mock: MockRouter) -> None: ) as stream: await stream.until_done() + @pytest.mark.asyncio + @pytest.mark.respx(base_url=base_url) + async def test_tool_use_long_string_input(self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: + respx_mock.post("/v1/messages").mock( + return_value=httpx2.Response(200, content=to_async_iter(get_long_tool_input_response())) + ) + parsed_sizes = record_parsed_json_sizes(monkeypatch) + + async with async_client.messages.stream( + max_tokens=1024, + messages=[{"role": "user", "content": "Write the notes file"}], + model="claude-sonnet-5-5", + ) as stream: + snapshots = [event.snapshot async for event in stream if event.type == "input_json"] + tool_use = (await stream.get_final_message()).content[0] + + assert tool_use.type == "tool_use" + assert_long_tool_input_parsed(snapshots, tool_use.input, parsed_sizes) + @pytest.mark.asyncio @pytest.mark.respx(base_url=base_url) async def test_tool_use_caller_omitted(self, respx_mock: MockRouter) -> None: