Prefer tool roots in streamed snapshots
Keep the streaming early-stop caller untouched by making extract_json_root_string prefer the first complete JSON root that normalizes as a tool request while preserving its first-root fallback for non-tool JSON. Add regression coverage for incidental JSON before a streamed tool-call envelope.
Alessandro committed
Jun 25, 2026 at 13:21 UTC
53a12a2ddd70afd72d379cc2245b63ed2ba0d090
3 files changed
+106
-8
helpers/extract_tools.py
+21
-8
@@ -10,12 +10,9 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None:
10
11
parsed_candidates: list[dict[str, Any]] = []
12
for ext_json in extract_json_root_strings(json.strip()):
13
- try:
14
- data = DirtyJson.parse_string(ext_json)
15
- if isinstance(data, dict):
16
- parsed_candidates.append(data)
17
- except Exception:
18
- continue
13
+ data = _parse_json_root_object(ext_json)
14
+ if data is not None:
15
+ parsed_candidates.append(data)
16
for data in parsed_candidates:
17
try:
18
normalize_tool_request(data)
@@ -53,9 +50,17 @@ def normalize_tool_request(tool_request: Any) -> tuple[str, dict]:
50
51
52
def extract_json_root_string(content: str) -> str | None:
56
- for root in extract_json_root_strings(content):
53
+ roots = extract_json_root_strings(content)
54
+ for root in roots:
55
+ data = _parse_json_root_object(root)
56
+ if data is None:
57
+ continue
58
+ try:
59
+ normalize_tool_request(data)
60
+ except ValueError:
61
+ continue
62
return root
58
- return None
63
+ return roots[0] if roots else None
64
65
66
def extract_json_root_strings(content: str) -> list[str]:
@@ -83,6 +88,14 @@ def extract_json_root_strings(content: str) -> list[str]:
88
return roots
89
90
91
+def _parse_json_root_object(root: str) -> dict[str, Any] | None:
92
+ try:
93
+ data = DirtyJson.parse_string(root)
94
+ except Exception:
95
+ return None
96
+ return data if isinstance(data, dict) else None
97
+
98
+
99
def extract_json_object_string(content):
100
start = content.find("{")
101
if start == -1:
helpers/extract_tools.py.dox.md
+1
@@ -25,6 +25,7 @@
25
- Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.
26
- Observed side-effect areas: settings/state persistence.
27
- Dirty parsing scans complete JSON object roots in prose and prefers the first object that normalizes as a valid tool request, so a leading text preamble or incidental non-tool object does not force a misformat warning when a valid tool call follows.
28
+- Streaming tool snapshots use the same valid-tool preference through `extract_json_root_string`, while preserving the first complete object fallback when no valid tool-call object is present.
29
- Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`.
30
31
## Key Concepts
tests/test_stream_tool_early_stop.py
+84
@@ -90,6 +90,20 @@ def test_json_parse_dirty_prefers_valid_tool_request_after_preamble_object():
90
}
91
92
93
+def test_extract_json_root_string_prefers_valid_tool_request():
94
+ text = (
95
+ 'I will call the tool after this note {"note":"not the tool"}.\n'
96
+ '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
97
+ )
98
+
99
+ assert extract_tools.extract_json_root_string(text) == (
100
+ '{"tool_name":"response","tool_args":{"text":"ok"}}'
101
+ )
102
+ assert extract_tools.extract_json_root_string(
103
+ 'Only a note {"note":"not the tool"}'
104
+ ) == '{"note":"not the tool"}'
105
+
106
+
107
def test_litellm_global_kwargs_merge_defaults_and_config(monkeypatch):
108
monkeypatch.setattr(
109
models.settings,
@@ -242,6 +256,76 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):
256
assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'
257
258
259
+@pytest.mark.asyncio
260
+async def test_unified_call_stops_after_tool_root_with_incidental_json(monkeypatch):
261
+ stream = _AsyncChunkStream(
262
+ [
263
+ {"type": "response.created"},
264
+ _response_event('Preamble {"note":"not the tool"}.\n'),
265
+ _response_event(
266
+ '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
267
+ ),
268
+ _response_event(" unreachable"),
269
+ ]
270
+ )
271
+
272
+ async def fake_aresponses(*args, **kwargs):
273
+ assert kwargs["stream"] is True
274
+ assert kwargs["input"] == ""
275
+ assert kwargs["store"] is True
276
+ return stream
277
+
278
+ async def fake_rate_limiter(*args, **kwargs):
279
+ return None
280
+
281
+ monkeypatch.setattr(litellm_transport, "aresponses", fake_aresponses)
282
+ monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)
283
+ monkeypatch.setattr(
284
+ models.settings,
285
+ "get_settings",
286
+ lambda: {"litellm_global_kwargs": {}},
287
+ )
288
+
289
+ wrapper = models.LiteLLMChatWrapper(
290
+ model="test-model",
291
+ provider="openai",
292
+ model_config=None,
293
+ )
294
+
295
+ seen: list[tuple[str, str]] = []
296
+
297
+ async def response_callback(chunk: str, full: str):
298
+ seen.append((chunk, full))
299
+ snapshot = extract_tools.extract_json_root_string(full)
300
+ if not snapshot:
301
+ return None
302
+ parsed_snapshot = extract_tools.json_parse_dirty(snapshot)
303
+ if parsed_snapshot is None:
304
+ return None
305
+ try:
306
+ extract_tools.normalize_tool_request(parsed_snapshot)
307
+ except ValueError:
308
+ return None
309
+ return snapshot
310
+
311
+ response, reasoning = await wrapper.unified_call(
312
+ messages=[],
313
+ response_callback=response_callback,
314
+ )
315
+
316
+ assert response == '{"tool_name":"response","tool_args":{"text":"ok"}}'
317
+ assert reasoning == ""
318
+ assert stream.index == 3
319
+ assert stream.closed is True
320
+ assert len(seen) == 2
321
+ assert seen[0][1] == 'Preamble {"note":"not the tool"}.\n'
322
+ assert (
323
+ seen[1][1]
324
+ == 'Preamble {"note":"not the tool"}.\n'
325
+ '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
326
+ )
327
+
328
+
329
@pytest.mark.asyncio
330
async def test_unified_call_closes_responses_stream_when_callback_raises(monkeypatch):
331
stream = _AsyncChunkStream([_response_event("interrupt me")])