ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix safety net, DRAINING metadata, and test import path

Browse source 
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.
This commit is contained in:
Daniel Han 2026-03-27 08:22:24 +00:00
commit 32ce2324f0
2 changed files with 198 additions and 256 deletions
Download patch file Download diff file Expand all files Collapse all files

51
studio/backend/core/inference/llama_cpp.py
View file

@ -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

401
tests/test_streaming_tool_detection.py
View file

@ -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 <function= prefix and drain."""
backend = _make_backend()
content = "<function=web_search><parameter=query>hello world</parameter></function>"
content = '<function=web_search><parameter=query>hello world</parameter></function>'
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)}
sse = _build_sse_stream(chunks, final_usage = usage)
sse = _build_sse_stream(chunks, final_usage=usage)
events = _collect_events(backend, sse)
@ -356,41 +321,48 @@ def test_whitespace_before_tool_xml():
rest = f"<tool_call>{tc_json}</tool_call>"
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"<tool_call>{tc_json}</tool_call>"}))
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 <think>...</think> 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 <think> tag
assert (
"<think>" in content_events[0]["text"]
), "First content should have <think> tag"
# Last content should have the actual answer
# Content should contain both <think> tags and the answer
final = content_events[-1]["text"]
assert "42" in final, f"Final content should have answer: {final}"
assert "<think>" in final, f"Should have <think> tag: {final}"
assert "Let me think" in final, f"Should have reasoning: {final}"
assert "42" in final, f"Should have answer: {final}"
assert "</think>" in final, f"Should have closing </think>: {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 "<think>" 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 <think> 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 <think> wrapper (reasoning-only fallback)
assert "<think>" not in final, f"Reasoning-only should not have <think> 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": "<tool"}),
_make_chunk({"content": "_tip>Use a screwdriver</tool_tip>"}),
_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()