diff --git a/helpers/extract_tools.py b/helpers/extract_tools.py index 188bf7fef..2e7d9eb6d 100644 --- a/helpers/extract_tools.py +++ b/helpers/extract_tools.py @@ -10,12 +10,9 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None: parsed_candidates: list[dict[str, Any]] = [] for ext_json in extract_json_root_strings(json.strip()): - try: - data = DirtyJson.parse_string(ext_json) - if isinstance(data, dict): - parsed_candidates.append(data) - except Exception: - continue + data = _parse_json_root_object(ext_json) + if data is not None: + parsed_candidates.append(data) for data in parsed_candidates: try: normalize_tool_request(data) @@ -53,9 +50,17 @@ def normalize_tool_request(tool_request: Any) -> tuple[str, dict]: def extract_json_root_string(content: str) -> str | None: - for root in extract_json_root_strings(content): + roots = extract_json_root_strings(content) + for root in roots: + data = _parse_json_root_object(root) + if data is None: + continue + try: + normalize_tool_request(data) + except ValueError: + continue return root - return None + return roots[0] if roots else None def extract_json_root_strings(content: str) -> list[str]: @@ -83,6 +88,14 @@ def extract_json_root_strings(content: str) -> list[str]: return roots +def _parse_json_root_object(root: str) -> dict[str, Any] | None: + try: + data = DirtyJson.parse_string(root) + except Exception: + return None + return data if isinstance(data, dict) else None + + def extract_json_object_string(content): start = content.find("{") if start == -1: diff --git a/helpers/extract_tools.py.dox.md b/helpers/extract_tools.py.dox.md index 6efb0de7f..0d5f7fdd1 100644 --- a/helpers/extract_tools.py.dox.md +++ b/helpers/extract_tools.py.dox.md @@ -25,6 +25,7 @@ - Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change. - Observed side-effect areas: settings/state persistence. - 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. +- 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. - Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`. ## Key Concepts diff --git a/tests/test_stream_tool_early_stop.py b/tests/test_stream_tool_early_stop.py index 0cb7a526d..791e70e78 100644 --- a/tests/test_stream_tool_early_stop.py +++ b/tests/test_stream_tool_early_stop.py @@ -90,6 +90,20 @@ def test_json_parse_dirty_prefers_valid_tool_request_after_preamble_object(): } +def test_extract_json_root_string_prefers_valid_tool_request(): + text = ( + 'I will call the tool after this note {"note":"not the tool"}.\n' + '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text' + ) + + assert extract_tools.extract_json_root_string(text) == ( + '{"tool_name":"response","tool_args":{"text":"ok"}}' + ) + assert extract_tools.extract_json_root_string( + 'Only a note {"note":"not the tool"}' + ) == '{"note":"not the tool"}' + + def test_litellm_global_kwargs_merge_defaults_and_config(monkeypatch): monkeypatch.setattr( models.settings, @@ -242,6 +256,76 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch): assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text' +@pytest.mark.asyncio +async def test_unified_call_stops_after_tool_root_with_incidental_json(monkeypatch): + stream = _AsyncChunkStream( + [ + {"type": "response.created"}, + _response_event('Preamble {"note":"not the tool"}.\n'), + _response_event( + '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text' + ), + _response_event(" unreachable"), + ] + ) + + async def fake_aresponses(*args, **kwargs): + assert kwargs["stream"] is True + assert kwargs["input"] == "" + assert kwargs["store"] is True + return stream + + async def fake_rate_limiter(*args, **kwargs): + return None + + monkeypatch.setattr(litellm_transport, "aresponses", fake_aresponses) + monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter) + monkeypatch.setattr( + models.settings, + "get_settings", + lambda: {"litellm_global_kwargs": {}}, + ) + + wrapper = models.LiteLLMChatWrapper( + model="test-model", + provider="openai", + model_config=None, + ) + + seen: list[tuple[str, str]] = [] + + async def response_callback(chunk: str, full: str): + seen.append((chunk, full)) + snapshot = extract_tools.extract_json_root_string(full) + if not snapshot: + return None + parsed_snapshot = extract_tools.json_parse_dirty(snapshot) + if parsed_snapshot is None: + return None + try: + extract_tools.normalize_tool_request(parsed_snapshot) + except ValueError: + return None + return snapshot + + response, reasoning = await wrapper.unified_call( + messages=[], + response_callback=response_callback, + ) + + assert response == '{"tool_name":"response","tool_args":{"text":"ok"}}' + assert reasoning == "" + assert stream.index == 3 + assert stream.closed is True + assert len(seen) == 2 + assert seen[0][1] == 'Preamble {"note":"not the tool"}.\n' + assert ( + seen[1][1] + == 'Preamble {"note":"not the tool"}.\n' + '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text' + ) + + @pytest.mark.asyncio async def test_unified_call_closes_responses_stream_when_callback_raises(monkeypatch): stream = _AsyncChunkStream([_response_event("interrupt me")])