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., 'Hello
').
- Buffer should hold briefly then flush when no prefix match at 32 chars."""
- backend = _make_backend()
-
- content = "This is a paragraph of HTML content that is not a tool call
"
- 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 "" 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 'Use a screwdriver"}),
- _make_chunk({}, finish_reason = "stop"),
- ]
- sse = _build_sse_stream(chunks)
- events = _collect_events(backend, sse)
-
- # "" so it enters BUFFERING.
- # Then "_tip>" does NOT match "" since the buffer becomes
- # "..." which doesn't start with "" or "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_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"{tc_json}"
-
- 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}")