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()