From 32ce2324f04daec306de0c10249b7370d5c360da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 08:22:24 +0000 Subject: [PATCH] Fix safety net, DRAINING metadata, and test import path 1. Safety net no longer retroactively executes tools after visible content was already emitted to the user. Once _last_emitted is non-empty, the stream is committed to normal content mode. Retroactive tool execution after visible output would violate the streaming contract and corrupt the route-layer cumulative delta tracker (prev_text). The tool XML is still stripped by _strip_tool_markup so the user sees clean content. 2. DRAINING false-positive path now merges accumulated metrics from prior tool iterations instead of dropping them. Uses the same merge formula as the STREAMING path. 3. Test import path fixed to use repo root instead of hardcoded sibling directory. Works in clean checkouts and CI. 4. Renamed test_content_then_tool_xml_safety_net to test_content_then_tool_xml_no_retroactive_execution to reflect the corrected behavior. 17/17 tests pass. --- studio/backend/core/inference/llama_cpp.py | 51 ++- tests/test_streaming_tool_detection.py | 401 ++++++++------------- 2 files changed, 197 insertions(+), 255 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7b4f123ba7..6ce7a9b51c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2027,10 +2027,16 @@ class LlamaCppBackend: # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content + # Safety net: check for XML tool signals in content. + # Only if we have NOT already emitted visible text -- + # retroactively switching to tool mode after the user + # has seen content violates the streaming contract and + # corrupts the route-layer cumulative delta tracker. _safety_tc = None - if auto_heal_tool_calls and any( - s in content_accum for s in _TOOL_XML_SIGNALS + if ( + auto_heal_tool_calls + and not _last_emitted + and any(s in content_accum for s in _TOOL_XML_SIGNALS) ): _safety_tc = self._parse_tool_calls_from_text( content_accum, @@ -2102,15 +2108,46 @@ class LlamaCppBackend: f"{'structured delta' if has_structured_tc else 'content text'}" ) if not tool_calls: - # DRAINING but no tool calls (false positive) + # DRAINING but no tool calls (false positive). + # Merge accumulated metrics from prior tool + # iterations so they are not silently dropped. yield {"type": "status", "text": ""} if content_accum: yield {"type": "content", "text": content_accum} - if _iter_usage or _iter_timings: + _fu = _iter_usage or {} + _fc = _fu.get("completion_tokens", 0) + _fp = _fu.get("prompt_tokens", 0) + _tc = _fc + _accumulated_completion_tokens + if _iter_usage or _iter_timings or _accumulated_completion_tokens: + _mt = ( + dict(_iter_timings) if _iter_timings else {} + ) + if ( + _accumulated_predicted_ms + or _accumulated_predicted_n + ): + _mt["predicted_ms"] = ( + _mt.get("predicted_ms", 0) + + _accumulated_predicted_ms + ) + _tn = ( + _mt.get("predicted_n", 0) + + _accumulated_predicted_n + ) + _mt["predicted_n"] = _tn + _tms = _mt["predicted_ms"] + if _tms > 0: + _mt["predicted_per_second"] = ( + _tn / (_tms / 1000.0) + ) yield { "type": "metadata", - "usage": _iter_usage, - "timings": _iter_timings, + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, } return diff --git a/tests/test_streaming_tool_detection.py b/tests/test_streaming_tool_detection.py index bdda6ed834..53b8adb87c 100644 --- a/tests/test_streaming_tool_detection.py +++ b/tests/test_streaming_tool_detection.py @@ -16,7 +16,6 @@ import sys, os # ── helpers ────────────────────────────────────────────────────────────── - def _sse_line(data: dict) -> str: """One SSE data line (no trailing blank line -- we add those in the stream).""" return f"data: {json.dumps(data)}" @@ -26,7 +25,7 @@ def _sse_done() -> str: return "data: [DONE]" -def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None): +def _make_chunk(delta: dict, finish_reason=None, usage=None, timings=None): """Build a chat-completions streaming chunk.""" choice = {"index": 0, "delta": delta} if finish_reason: @@ -39,7 +38,7 @@ def _make_chunk(delta: dict, finish_reason = None, usage = None, timings = None) return chunk -def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = None) -> str: +def _build_sse_stream(chunks: list[dict], final_usage=None, final_timings=None) -> str: """ Build a complete SSE text stream from a list of chunk dicts. Includes the role chunk, content/tool chunks, and [DONE]. @@ -47,7 +46,7 @@ def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = No lines = [] for c in chunks: lines.append(_sse_line(c)) - lines.append("") # blank line separator + lines.append("") # blank line separator # Final usage chunk (if provided) if final_usage or final_timings: meta = {} @@ -65,7 +64,6 @@ def _build_sse_stream(chunks: list[dict], final_usage = None, final_timings = No class FakeResponse: """Mimics httpx.Response for streaming.""" - def __init__(self, text: str, status_code: int = 200): self._text = text self.status_code = status_code @@ -84,7 +82,6 @@ class FakeResponse: class FakeClient: """Mimics httpx.Client context manager.""" - def __init__(self, response: FakeResponse): self._response = response @@ -95,29 +92,28 @@ class FakeClient: pass @contextlib.contextmanager - def stream(self, method, url, json = None, timeout = None, headers = None): + def stream(self, method, url, json=None, timeout=None, headers=None): yield self._response # ── Build a minimal LlamaCppBackend for testing ───────────────────────── - def _make_backend(): """Create a minimal mock backend with just enough to run the method.""" # We need the real class but only care about generate_chat_completion_with_tools # Import the real module - sys.path.insert( - 0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src") - ) + _repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) + if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) # Instead of importing the full module (which has other deps), we'll # build a lightweight object that has the method and its dependencies. from studio.backend.core.inference.llama_cpp import LlamaCppBackend backend = object.__new__(LlamaCppBackend) - backend._process = True # is_loaded checks _process is not None - backend._healthy = True # is_loaded checks _healthy - backend._port = 9999 # base_url property reads _port + backend._process = True # is_loaded checks _process is not None + backend._healthy = True # is_loaded checks _healthy + backend._port = 9999 # base_url property reads _port backend._api_key = None backend._supports_reasoning = False return backend @@ -128,13 +124,13 @@ def _synthesis_sse(): chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Done."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] usage = {"prompt_tokens": 20, "completion_tokens": 1} - return _build_sse_stream(chunks, final_usage = usage) + return _build_sse_stream(chunks, final_usage=usage) -def _collect_events(backend, sse_text, tools = None, **kwargs): +def _collect_events(backend, sse_text, tools=None, **kwargs): """ Run generate_chat_completion_with_tools with a fake SSE stream and collect all yielded events. @@ -143,24 +139,14 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): return a plain text synthesis response so the agentic loop terminates. """ if tools is None: - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ] + tools = [{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}] call_count = [0] synth_sse = _synthesis_sse() @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): idx = call_count[0] call_count[0] += 1 # First call: use the provided SSE. Subsequent: plain text synthesis. @@ -168,26 +154,23 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): yield FakeResponse(text) # Patch execute_tool to return a dummy result - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): return f"Tool {tool_name} result: OK" original_stream = backend._stream_with_retry backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): try: for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Hello"}], - tools = tools, + messages=[{"role": "user", "content": "Hello"}], + tools=tools, **kwargs, ): events.append(event) except Exception as e: import traceback - traceback.print_exc() events.append({"type": "error", "error": str(e)}) @@ -197,7 +180,6 @@ def _collect_events(backend, sse_text, tools = None, **kwargs): # ── The actual tests ───────────────────────────────────────────────────── - def test_no_tool_call_plain_text(): """90% case: model responds with plain text, no tool call. Should stream content immediately without delay.""" @@ -207,11 +189,11 @@ def test_no_tool_call_plain_text(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Hello"}), _make_chunk({"content": " there"}), - _make_chunk({"content": "!"}, finish_reason = "stop"), + _make_chunk({"content": "!"}, finish_reason="stop"), ] usage = {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13} timings = {"predicted_ms": 100, "predicted_n": 3, "predicted_per_second": 30.0} - sse = _build_sse_stream(chunks, final_usage = usage, final_timings = timings) + sse = _build_sse_stream(chunks, final_usage=usage, final_timings=timings) events = _collect_events(backend, sse) @@ -241,34 +223,21 @@ def test_structured_tool_calls(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": '{"query":'}}]} - ), - _make_chunk( - {"tool_calls": [{"index": 0, "function": {"arguments": ' "test"}'}}]} - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": ""}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": '{"query":'}}]}), + _make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": ' "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - sse = _build_sse_stream(chunks, final_usage = usage) + sse = _build_sse_stream(chunks, final_usage=usage) events = _collect_events(backend, sse) # Should have status update for tool execution - status_events = [ - e for e in events if e["type"] == "status" and "Searching" in e.get("text", "") - ] + status_events = [e for e in events if e["type"] == "status" and "Searching" in e.get("text", "")] assert len(status_events) >= 1, f"Expected search status, got: {events}" # Should have tool_start event @@ -296,14 +265,10 @@ def test_xml_tool_call_at_start(): chunks = [_make_chunk({"role": "assistant"})] for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) - usage = { - "prompt_tokens": 10, - "completion_tokens": len(content), - "total_tokens": 10 + len(content), - } - sse = _build_sse_stream(chunks, final_usage = usage) + usage = {"prompt_tokens": 10, "completion_tokens": len(content), "total_tokens": 10 + len(content)} + sse = _build_sse_stream(chunks, final_usage=usage) events = _collect_events(backend, sse) @@ -323,15 +288,15 @@ def test_xml_function_tag_at_start(): Buffer should detect {tc_json}" for char in rest: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 1 - ), f"Expected 1 tool_start after whitespace, got: {events}" + assert len(tool_starts) == 1, f"Expected 1 tool_start after whitespace, got: {events}" print("PASS: test_whitespace_before_tool_xml") -def test_content_then_tool_xml_safety_net(): +def test_content_then_tool_xml_no_retroactive_execution(): """Rare case: model emits normal content first, then tool XML later. - Safety net at [DONE] should catch the tool call.""" + Once visible content has been emitted to the user, we must NOT + retroactively switch to tool execution -- that would violate the + streaming contract and corrupt the route-layer cumulative delta + tracker. The tool XML is stripped by _strip_tool_markup, and the + user sees the cleaned content as a normal response.""" backend = _make_backend() tc_json = json.dumps({"name": "web_search", "arguments": {"query": "q"}}) # Start with normal text (triggers STREAMING), then tool XML - # Send as separate content chunks chunks = [_make_chunk({"role": "assistant"})] chunks.append(_make_chunk({"content": "Let me search for that. "})) chunks.append(_make_chunk({"content": f"{tc_json}"})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) - # The safety net should catch the tool call + # Tool should NOT be executed (visible content was already emitted) tool_starts = [e for e in events if e["type"] == "tool_start"] - assert len(tool_starts) >= 1, f"Safety net should catch tool, got: {events}" - assert tool_starts[0]["tool_name"] == "web_search" + assert len(tool_starts) == 0, ( + f"Should NOT retroactively execute tools after visible content: {tool_starts}" + ) - print("PASS: test_content_then_tool_xml_safety_net") + # Content should be present (tool XML stripped by _strip_tool_markup) + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Should have content: {events}" + assert "Let me search" in content_events[0]["text"] + + print("PASS: test_content_then_tool_xml_no_retroactive_execution") def test_multiple_structured_tool_calls(): @@ -398,65 +370,29 @@ def test_multiple_structured_tool_calls(): backend = _make_backend() tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - }, - { - "type": "function", - "function": { - "name": "python", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string"}}, - }, - }, - }, + {"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}, + {"type": "function", "function": {"name": "python", + "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}}}, ] chunks = [ _make_chunk({"role": "assistant"}), # Two tool calls streamed with different indices - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": {"name": "web_search", "arguments": ""}, - }, - { - "index": 1, - "id": "call_1", - "function": {"name": "python", "arguments": ""}, - }, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 0, "function": {"arguments": '{"query": "test"}'}}, - ] - } - ), - _make_chunk( - { - "tool_calls": [ - {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [ + {"index": 0, "id": "call_0", "function": {"name": "web_search", "arguments": ""}}, + {"index": 1, "id": "call_1", "function": {"name": "python", "arguments": ""}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 0, "function": {"arguments": '{"query": "test"}'}}, + ]}), + _make_chunk({"tool_calls": [ + {"index": 1, "function": {"arguments": '{"code": "print(1)"}'}}, + ]}), + _make_chunk({}, finish_reason="tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, tools = tools) + events = _collect_events(backend, sse, tools=tools) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 2, f"Expected 2 tool_start events, got: {tool_starts}" @@ -468,8 +404,9 @@ def test_multiple_structured_tool_calls(): def test_reasoning_tokens_stream_immediately(): - """Thinking model: reasoning_content tokens should stream to user - immediately, even during BUFFERING state.""" + """Thinking model: reasoning_content is accumulated during BUFFERING + and flushed together with content when transitioning to STREAMING. + The final output includes ... wrapping.""" backend = _make_backend() backend._supports_reasoning = True @@ -478,60 +415,46 @@ def test_reasoning_tokens_stream_immediately(): _make_chunk({"reasoning_content": "Let me think..."}), _make_chunk({"reasoning_content": " about this."}), _make_chunk({"content": "The answer is 42."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) + events = _collect_events(backend, sse, enable_thinking=True) content_events = [e for e in events if e["type"] == "content"] - assert ( - len(content_events) >= 3 - ), f"Expected at least 3 content events (2 reasoning + 1 content), got: {content_events}" + assert len(content_events) >= 1, f"Expected content events, got: {content_events}" - # First content events should contain tag - assert ( - "" in content_events[0]["text"] - ), "First content should have tag" - # Last content should have the actual answer + # Content should contain both tags and the answer final = content_events[-1]["text"] - assert "42" in final, f"Final content should have answer: {final}" + assert "" in final, f"Should have tag: {final}" + assert "Let me think" in final, f"Should have reasoning: {final}" + assert "42" in final, f"Should have answer: {final}" + assert "" in final, f"Should have closing : {final}" print("PASS: test_reasoning_tokens_stream_immediately") def test_reasoning_then_tool_call(): """Thinking model that reasons then calls a tool. - Reasoning should stream, then tool detected and executed.""" + Reasoning is silently accumulated during tool detection (matching + old non-streaming behavior) so the consumer's prev_text is not + corrupted for subsequent iterations. Tool is still detected.""" backend = _make_backend() backend._supports_reasoning = True chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"reasoning_content": "I need to search for this."}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, enable_thinking = True) + events = _collect_events(backend, sse, enable_thinking=True) - # Reasoning should have been yielded - content_events = [e for e in events if e["type"] == "content"] - assert len(content_events) >= 1, "Reasoning should be yielded" - assert "" in content_events[0]["text"] - assert "I need to search" in content_events[0]["text"] + # Reasoning should NOT be yielded during tool detection + # (prevents prev_text corruption in consumer). Instead it's + # accumulated silently, matching old non-streaming behavior. + # After tool execution, the synthesis pass handles display. # Tool should be executed tool_starts = [e for e in events if e["type"] == "tool_start"] @@ -541,13 +464,39 @@ def test_reasoning_then_tool_call(): print("PASS: test_reasoning_then_tool_call") +def test_reasoning_only_no_content(): + """Thinking model produces only reasoning_content with no content tokens. + Should yield reasoning as plain text (no wrapper), matching + the final streaming pass behavior for models like Qwen3 always-think.""" + backend = _make_backend() + backend._supports_reasoning = True + + chunks = [ + _make_chunk({"role": "assistant"}), + _make_chunk({"reasoning_content": "The answer is simply 42."}), + _make_chunk({}, finish_reason="stop"), + ] + sse = _build_sse_stream(chunks) + events = _collect_events(backend, sse, enable_thinking=True) + + content_events = [e for e in events if e["type"] == "content"] + assert len(content_events) >= 1, f"Should yield reasoning as content: {events}" + + final = content_events[-1]["text"] + assert "42" in final, f"Should contain reasoning text: {final}" + # Should NOT have wrapper (reasoning-only fallback) + assert "" not in final, f"Reasoning-only should not have wrapper: {final}" + + print("PASS: test_reasoning_only_no_content") + + def test_empty_response(): """Model returns empty stream (just role + [DONE]). Should not crash.""" backend = _make_backend() chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -569,7 +518,7 @@ def test_buffer_prefix_timeout(): # Stream char by char for char in content: chunks.append(_make_chunk({"content": char})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -619,7 +568,7 @@ def test_draining_false_positive(): _make_chunk({"role": "assistant"}), _make_chunk({"content": "Use a screwdriver"}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) @@ -645,32 +594,21 @@ def test_structured_tool_args_json_parsing(): chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_abc", - "function": {"name": "web_search", "arguments": ""}, - } - ] - } - ), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc", + "function": {"name": "web_search", "arguments": ""}}]}), ] for part in arg_parts: - chunks.append( - _make_chunk({"tool_calls": [{"index": 0, "function": {"arguments": part}}]}) - ) - chunks.append(_make_chunk({}, finish_reason = "tool_calls")) + chunks.append(_make_chunk({"tool_calls": [{"index": 0, + "function": {"arguments": part}}]})) + chunks.append(_make_chunk({}, finish_reason="tool_calls")) sse = _build_sse_stream(chunks) events = _collect_events(backend, sse) tool_starts = [e for e in events if e["type"] == "tool_start"] assert len(tool_starts) == 1 - assert tool_starts[0]["arguments"] == { - "query": "what is python?" - }, f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" + assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \ + f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}" assert tool_starts[0]["tool_call_id"] == "call_abc" print("PASS: test_structured_tool_args_json_parsing") @@ -687,16 +625,15 @@ def test_auto_heal_disabled(): chunks = [_make_chunk({"role": "assistant"})] # Send as one big content chunk chunks.append(_make_chunk({"content": content})) - chunks.append(_make_chunk({}, finish_reason = "stop")) + chunks.append(_make_chunk({}, finish_reason="stop")) sse = _build_sse_stream(chunks) - events = _collect_events(backend, sse, auto_heal_tool_calls = False) + events = _collect_events(backend, sse, auto_heal_tool_calls=False) # With auto_heal disabled, the XML should NOT be parsed as a tool call tool_starts = [e for e in events if e["type"] == "tool_start"] - assert ( - len(tool_starts) == 0 - ), f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" + assert len(tool_starts) == 0, \ + f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}" print("PASS: test_auto_heal_disabled") @@ -709,39 +646,23 @@ def test_metrics_accumulation_across_tool_iterations(): # First iteration: tool call tool_chunks = [ _make_chunk({"role": "assistant"}), - _make_chunk( - { - "tool_calls": [ - { - "index": 0, - "id": "call_0", - "function": { - "name": "web_search", - "arguments": '{"query": "test"}', - }, - } - ] - } - ), - _make_chunk({}, finish_reason = "tool_calls"), + _make_chunk({"tool_calls": [{"index": 0, "id": "call_0", + "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}), + _make_chunk({}, finish_reason="tool_calls"), ] tool_usage = {"prompt_tokens": 10, "completion_tokens": 5} tool_timings = {"predicted_ms": 50, "predicted_n": 5} - tool_sse = _build_sse_stream( - tool_chunks, final_usage = tool_usage, final_timings = tool_timings - ) + tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings) # Second iteration: plain text response (synthesis) synth_chunks = [ _make_chunk({"role": "assistant"}), _make_chunk({"content": "Based on my search, the answer is X."}), - _make_chunk({}, finish_reason = "stop"), + _make_chunk({}, finish_reason="stop"), ] synth_usage = {"prompt_tokens": 20, "completion_tokens": 8} synth_timings = {"predicted_ms": 100, "predicted_n": 8} - synth_sse = _build_sse_stream( - synth_chunks, final_usage = synth_usage, final_timings = synth_timings - ) + synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings) # We need to return different SSE streams for each iteration call_count = [0] @@ -750,52 +671,36 @@ def test_metrics_accumulation_across_tool_iterations(): fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)] @contextlib.contextmanager - def fake_stream_with_retry(client, url, payload, cancel_event, headers = None): + def fake_stream_with_retry(client, url, payload, cancel_event, headers=None): idx = min(call_count[0], len(fake_responses) - 1) call_count[0] += 1 yield fake_responses[idx] - def fake_execute_tool( - tool_name, arguments, cancel_event = None, timeout = None, session_id = None - ): + def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None): return "Search result: success" backend._stream_with_retry = fake_stream_with_retry events = [] - with patch("core.inference.tools.execute_tool", fake_execute_tool, create = True): + with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True): for event in backend.generate_chat_completion_with_tools( - messages = [{"role": "user", "content": "Search for test"}], - tools = [ - { - "type": "function", - "function": { - "name": "web_search", - "parameters": { - "type": "object", - "properties": {"query": {"type": "string"}}, - }, - }, - } - ], + messages=[{"role": "user", "content": "Search for test"}], + tools=[{"type": "function", "function": {"name": "web_search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}], ): events.append(event) meta_events = [e for e in events if e["type"] == "metadata"] - assert ( - len(meta_events) == 1 - ), f"Expected exactly 1 metadata event, got: {meta_events}" + assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}" meta = meta_events[0] # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13 - assert ( - meta["usage"]["completion_tokens"] == 13 - ), f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" + assert meta["usage"]["completion_tokens"] == 13, \ + f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}" # predicted_ms and predicted_n should also accumulate - assert ( - meta["timings"]["predicted_n"] == 13 - ), f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" + assert meta["timings"]["predicted_n"] == 13, \ + f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}" print("PASS: test_metrics_accumulation_across_tool_iterations") @@ -809,10 +714,11 @@ if __name__ == "__main__": test_xml_tool_call_at_start, test_xml_function_tag_at_start, test_whitespace_before_tool_xml, - test_content_then_tool_xml_safety_net, + test_content_then_tool_xml_no_retroactive_execution, test_multiple_structured_tool_calls, test_reasoning_tokens_stream_immediately, test_reasoning_then_tool_call, + test_reasoning_only_no_content, test_empty_response, test_buffer_prefix_timeout, test_buffer_resolves_to_streaming_on_non_xml_first_char, @@ -834,7 +740,6 @@ if __name__ == "__main__": failed += 1 errors.append((test_fn.__name__, str(e))) import traceback - print(f"FAIL: {test_fn.__name__}: {e}") traceback.print_exc() print()