848 lines
29 KiB
Python
848 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
|
|
sys.path.insert(
|
|
0, os.path.join(os.path.dirname(__file__), "..", "unsloth_studio_src")
|
|
)
|
|
|
|
# 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_safety_net():
|
|
"""Rare case: model emits normal content first, then tool XML later.
|
|
Safety net at [DONE] should catch the tool call."""
|
|
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"))
|
|
|
|
sse = _build_sse_stream(chunks)
|
|
events = _collect_events(backend, sse)
|
|
|
|
# The safety net should catch the tool call
|
|
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"
|
|
|
|
print("PASS: test_content_then_tool_xml_safety_net")
|
|
|
|
|
|
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 tokens should stream to user
|
|
immediately, even during BUFFERING state."""
|
|
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) >= 3
|
|
), f"Expected at least 3 content events (2 reasoning + 1 content), 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
|
|
final = content_events[-1]["text"]
|
|
assert "42" in final, f"Final content should have answer: {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."""
|
|
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 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"]
|
|
|
|
# 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_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_safety_net,
|
|
test_multiple_structured_tool_calls,
|
|
test_reasoning_tokens_stream_immediately,
|
|
test_reasoning_then_tool_call,
|
|
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}")
|