32ce2324f0
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.
753 lines
29 KiB
Python
753 lines
29 KiB
Python
"""
|
|
Exhaustive tests for the speculative-buffer streaming tool detection in
|
|
generate_chat_completion_with_tools().
|
|
|
|
We mock the HTTP layer so llama-server is not required. Each test constructs
|
|
the exact SSE byte stream that llama-server would emit, feeds it through the
|
|
real method, and asserts on the yielded events.
|
|
"""
|
|
|
|
import json
|
|
import threading
|
|
import types
|
|
import contextlib
|
|
from unittest.mock import MagicMock, patch, PropertyMock
|
|
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)}"
|
|
|
|
|
|
def _sse_done() -> str:
|
|
return "data: [DONE]"
|
|
|
|
|
|
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:
|
|
choice["finish_reason"] = finish_reason
|
|
chunk = {"choices": [choice]}
|
|
if usage:
|
|
chunk["usage"] = usage
|
|
if timings:
|
|
chunk["timings"] = timings
|
|
return chunk
|
|
|
|
|
|
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].
|
|
"""
|
|
lines = []
|
|
for c in chunks:
|
|
lines.append(_sse_line(c))
|
|
lines.append("") # blank line separator
|
|
# Final usage chunk (if provided)
|
|
if final_usage or final_timings:
|
|
meta = {}
|
|
if final_usage:
|
|
meta["usage"] = final_usage
|
|
if final_timings:
|
|
meta["timings"] = final_timings
|
|
meta["choices"] = []
|
|
lines.append(_sse_line(meta))
|
|
lines.append("")
|
|
lines.append(_sse_done())
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
class FakeResponse:
|
|
"""Mimics httpx.Response for streaming."""
|
|
def __init__(self, text: str, status_code: int = 200):
|
|
self._text = text
|
|
self.status_code = status_code
|
|
self._closed = False
|
|
|
|
def iter_text(self):
|
|
# Yield the whole thing in one shot (simplest case)
|
|
yield self._text
|
|
|
|
def read(self):
|
|
return self._text.encode()
|
|
|
|
def close(self):
|
|
self._closed = True
|
|
|
|
|
|
class FakeClient:
|
|
"""Mimics httpx.Client context manager."""
|
|
def __init__(self, response: FakeResponse):
|
|
self._response = response
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
@contextlib.contextmanager
|
|
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
|
|
_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._api_key = None
|
|
backend._supports_reasoning = False
|
|
return backend
|
|
|
|
|
|
def _synthesis_sse():
|
|
"""Build a simple text SSE response for post-tool synthesis."""
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_make_chunk({"content": "Done."}),
|
|
_make_chunk({}, finish_reason="stop"),
|
|
]
|
|
usage = {"prompt_tokens": 20, "completion_tokens": 1}
|
|
return _build_sse_stream(chunks, final_usage=usage)
|
|
|
|
|
|
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.
|
|
|
|
After the first iteration (tool detection), subsequent iterations
|
|
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"}}}}}]
|
|
|
|
call_count = [0]
|
|
synth_sse = _synthesis_sse()
|
|
|
|
@contextlib.contextmanager
|
|
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.
|
|
text = sse_text if idx == 0 else synth_sse
|
|
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):
|
|
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):
|
|
try:
|
|
for event in backend.generate_chat_completion_with_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)})
|
|
|
|
backend._stream_with_retry = original_stream
|
|
return events
|
|
|
|
|
|
# ── 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."""
|
|
backend = _make_backend()
|
|
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_make_chunk({"content": "Hello"}),
|
|
_make_chunk({"content": " there"}),
|
|
_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)
|
|
|
|
events = _collect_events(backend, sse)
|
|
|
|
# Should have content events with cumulative text
|
|
content_events = [e for e in events if e["type"] == "content"]
|
|
assert len(content_events) >= 1, f"Expected content events, got: {events}"
|
|
|
|
# Final content should contain the full text
|
|
final_content = content_events[-1]["text"]
|
|
assert "Hello there!" in final_content, f"Missing text in: {final_content}"
|
|
|
|
# Should have metadata
|
|
meta_events = [e for e in events if e["type"] == "metadata"]
|
|
assert len(meta_events) == 1, f"Expected 1 metadata event, got: {meta_events}"
|
|
|
|
# Should have status clear
|
|
status_events = [e for e in events if e["type"] == "status"]
|
|
assert any(e["text"] == "" for e in status_events), "Missing status clear"
|
|
|
|
print("PASS: test_no_tool_call_plain_text")
|
|
|
|
|
|
def test_structured_tool_calls():
|
|
"""Model emits structured delta.tool_calls (the standard path).
|
|
Should detect instantly and execute."""
|
|
backend = _make_backend()
|
|
|
|
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"),
|
|
]
|
|
usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
|
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", "")]
|
|
assert len(status_events) >= 1, f"Expected search status, got: {events}"
|
|
|
|
# Should have tool_start event
|
|
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
|
assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {tool_starts}"
|
|
assert tool_starts[0]["tool_name"] == "web_search"
|
|
assert tool_starts[0]["arguments"] == {"query": "test"}
|
|
|
|
# Should have tool_end event
|
|
tool_ends = [e for e in events if e["type"] == "tool_end"]
|
|
assert len(tool_ends) == 1, f"Expected 1 tool_end, got: {tool_ends}"
|
|
|
|
print("PASS: test_structured_tool_calls")
|
|
|
|
|
|
def test_xml_tool_call_at_start():
|
|
"""Model emits <tool_call>JSON</tool_call> instead of structured tool_calls.
|
|
Buffer should detect <tool_call> prefix and drain."""
|
|
backend = _make_backend()
|
|
|
|
tc_json = json.dumps({"name": "web_search", "arguments": {"query": "hello"}})
|
|
content = f"<tool_call>{tc_json}</tool_call>"
|
|
|
|
# Stream the XML content token by token to simulate real streaming
|
|
chunks = [_make_chunk({"role": "assistant"})]
|
|
for char in content:
|
|
chunks.append(_make_chunk({"content": char}))
|
|
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)
|
|
|
|
events = _collect_events(backend, sse)
|
|
|
|
# Should detect tool call and execute it
|
|
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
|
assert len(tool_starts) == 1, f"Expected 1 tool_start, got: {events}"
|
|
assert tool_starts[0]["tool_name"] == "web_search"
|
|
|
|
tool_ends = [e for e in events if e["type"] == "tool_end"]
|
|
assert len(tool_ends) == 1, f"Expected 1 tool_end"
|
|
|
|
print("PASS: test_xml_tool_call_at_start")
|
|
|
|
|
|
def test_xml_function_tag_at_start():
|
|
"""Model emits <function=web_search> tag.
|
|
Buffer should detect <function= prefix and drain."""
|
|
backend = _make_backend()
|
|
|
|
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"))
|
|
|
|
usage = {"prompt_tokens": 10, "completion_tokens": len(content)}
|
|
sse = _build_sse_stream(chunks, final_usage=usage)
|
|
|
|
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, got: {events}"
|
|
assert tool_starts[0]["tool_name"] == "web_search"
|
|
|
|
print("PASS: test_xml_function_tag_at_start")
|
|
|
|
|
|
def test_whitespace_before_tool_xml():
|
|
"""Model emits whitespace then <tool_call>. Buffer should strip
|
|
leading whitespace before prefix check."""
|
|
backend = _make_backend()
|
|
|
|
tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}})
|
|
content = f" \n <tool_call>{tc_json}</tool_call>"
|
|
|
|
chunks = [_make_chunk({"role": "assistant"})]
|
|
# Send whitespace as one chunk, then the rest
|
|
chunks.append(_make_chunk({"content": " \n "}))
|
|
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"))
|
|
|
|
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}"
|
|
|
|
print("PASS: test_whitespace_before_tool_xml")
|
|
|
|
|
|
def test_content_then_tool_xml_no_retroactive_execution():
|
|
"""Rare case: model emits normal content first, then tool XML later.
|
|
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
|
|
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"))
|
|
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
# 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) == 0, (
|
|
f"Should NOT retroactively execute tools after visible content: {tool_starts}"
|
|
)
|
|
|
|
# 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():
|
|
"""Model calls two tools in one response (parallel 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"}}}}},
|
|
]
|
|
|
|
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"),
|
|
]
|
|
sse = _build_sse_stream(chunks)
|
|
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}"
|
|
|
|
names = {ts["tool_name"] for ts in tool_starts}
|
|
assert names == {"web_search", "python"}, f"Wrong tool names: {names}"
|
|
|
|
print("PASS: test_multiple_structured_tool_calls")
|
|
|
|
|
|
def test_reasoning_tokens_stream_immediately():
|
|
"""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
|
|
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_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"),
|
|
]
|
|
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"Expected content events, got: {content_events}"
|
|
|
|
# Content should contain both <think> tags and the answer
|
|
final = content_events[-1]["text"]
|
|
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 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"),
|
|
]
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse, enable_thinking=True)
|
|
|
|
# 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"]
|
|
assert len(tool_starts) == 1, f"Expected tool_start: {events}"
|
|
assert tool_starts[0]["tool_name"] == "web_search"
|
|
|
|
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"),
|
|
]
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
# Should not crash, just return with no content
|
|
error_events = [e for e in events if e.get("type") == "error"]
|
|
assert len(error_events) == 0, f"Should not error: {error_events}"
|
|
|
|
print("PASS: test_empty_response")
|
|
|
|
|
|
def test_buffer_prefix_timeout():
|
|
"""Content starts with '<' but is not a tool call (e.g., '<p>Hello</p>').
|
|
Buffer should hold briefly then flush when no prefix match at 32 chars."""
|
|
backend = _make_backend()
|
|
|
|
content = "<p>This is a paragraph of HTML content that is not a tool call</p>"
|
|
chunks = [_make_chunk({"role": "assistant"})]
|
|
# Stream char by char
|
|
for char in content:
|
|
chunks.append(_make_chunk({"content": char}))
|
|
chunks.append(_make_chunk({}, finish_reason="stop"))
|
|
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
content_events = [e for e in events if e["type"] == "content"]
|
|
assert len(content_events) >= 1, f"Should have content events: {events}"
|
|
|
|
# No tool calls should be detected
|
|
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
|
assert len(tool_starts) == 0, f"Should not detect tools in HTML: {tool_starts}"
|
|
|
|
# Final content should contain the HTML
|
|
final = content_events[-1]["text"]
|
|
assert "<p>" in final, f"HTML content should pass through: {final}"
|
|
|
|
print("PASS: test_buffer_prefix_timeout")
|
|
|
|
|
|
def test_buffer_resolves_to_streaming_on_non_xml_first_char():
|
|
"""First content char is not '<' and not whitespace.
|
|
Should immediately transition to STREAMING."""
|
|
backend = _make_backend()
|
|
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_make_chunk({"content": "H"}), # 'H' is not '<', instant STREAMING
|
|
_make_chunk({"content": "ello"}),
|
|
]
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
content_events = [e for e in events if e["type"] == "content"]
|
|
# First content event should appear immediately with just "H"
|
|
assert len(content_events) >= 1
|
|
assert "H" in content_events[0]["text"]
|
|
|
|
print("PASS: test_buffer_resolves_to_streaming_on_non_xml_first_char")
|
|
|
|
|
|
def test_draining_false_positive():
|
|
"""Buffer detects '<tool' prefix but stream ends before completing
|
|
the tag (e.g., '<tool' then EOF). Should yield content as-is."""
|
|
backend = _make_backend()
|
|
|
|
# Content that starts like a tool tag but isn't
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_make_chunk({"content": "<tool"}),
|
|
_make_chunk({"content": "_tip>Use a screwdriver</tool_tip>"}),
|
|
_make_chunk({}, finish_reason="stop"),
|
|
]
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
# "<tool" is a prefix of "<tool_call>" so it enters BUFFERING.
|
|
# Then "_tip>" does NOT match "<tool_call>" since the buffer becomes
|
|
# "<tool_tip>..." which doesn't start with "<tool_call>" or "<function=".
|
|
# But at >32 chars the buffer should flush.
|
|
# No tool should be executed.
|
|
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
|
assert len(tool_starts) == 0, f"Should not detect tool in <tool_tip>: {tool_starts}"
|
|
|
|
print("PASS: test_draining_false_positive")
|
|
|
|
|
|
def test_structured_tool_args_json_parsing():
|
|
"""Verify that arguments streamed across multiple chunks get reassembled
|
|
and parsed correctly as JSON."""
|
|
backend = _make_backend()
|
|
|
|
# Arguments split across 4 chunks
|
|
arg_parts = ['{"qu', 'ery":', ' "wha', 't is python?"}']
|
|
|
|
chunks = [
|
|
_make_chunk({"role": "assistant"}),
|
|
_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"))
|
|
|
|
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]["tool_call_id"] == "call_abc"
|
|
|
|
print("PASS: test_structured_tool_args_json_parsing")
|
|
|
|
|
|
def test_auto_heal_disabled():
|
|
"""When auto_heal_tool_calls=False, XML tool calls in content should NOT
|
|
be parsed -- only structured tool_calls are honored."""
|
|
backend = _make_backend()
|
|
|
|
tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}})
|
|
content = f"<tool_call>{tc_json}</tool_call>"
|
|
|
|
chunks = [_make_chunk({"role": "assistant"})]
|
|
# Send as one big content chunk
|
|
chunks.append(_make_chunk({"content": content}))
|
|
chunks.append(_make_chunk({}, finish_reason="stop"))
|
|
|
|
sse = _build_sse_stream(chunks)
|
|
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}"
|
|
|
|
print("PASS: test_auto_heal_disabled")
|
|
|
|
|
|
def test_metrics_accumulation_across_tool_iterations():
|
|
"""When tools are called, metrics from the tool iteration should be
|
|
accumulated and included in the final metadata."""
|
|
backend = _make_backend()
|
|
|
|
# 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"),
|
|
]
|
|
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)
|
|
|
|
# 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"),
|
|
]
|
|
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)
|
|
|
|
# We need to return different SSE streams for each iteration
|
|
call_count = [0]
|
|
original_sse = [tool_sse, synth_sse]
|
|
|
|
fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)]
|
|
|
|
@contextlib.contextmanager
|
|
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):
|
|
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):
|
|
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"}}}}}],
|
|
):
|
|
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}"
|
|
|
|
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']}"
|
|
|
|
# predicted_ms and predicted_n should also accumulate
|
|
assert meta["timings"]["predicted_n"] == 13, \
|
|
f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}"
|
|
|
|
print("PASS: test_metrics_accumulation_across_tool_iterations")
|
|
|
|
|
|
# ── Run all tests ────────────────────────────────────────────────────────
|
|
|
|
if __name__ == "__main__":
|
|
tests = [
|
|
test_no_tool_call_plain_text,
|
|
test_structured_tool_calls,
|
|
test_xml_tool_call_at_start,
|
|
test_xml_function_tag_at_start,
|
|
test_whitespace_before_tool_xml,
|
|
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,
|
|
test_draining_false_positive,
|
|
test_structured_tool_args_json_parsing,
|
|
test_auto_heal_disabled,
|
|
test_metrics_accumulation_across_tool_iterations,
|
|
]
|
|
|
|
passed = 0
|
|
failed = 0
|
|
errors = []
|
|
|
|
for test_fn in tests:
|
|
try:
|
|
test_fn()
|
|
passed += 1
|
|
except Exception as e:
|
|
failed += 1
|
|
errors.append((test_fn.__name__, str(e)))
|
|
import traceback
|
|
print(f"FAIL: {test_fn.__name__}: {e}")
|
|
traceback.print_exc()
|
|
print()
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Results: {passed} passed, {failed} failed, {len(tests)} total")
|
|
if errors:
|
|
print(f"\nFailed tests:")
|
|
for name, err in errors:
|
|
print(f" - {name}: {err}")
|
|
print(f"{'='*60}")
|