studio: deterministic backend tool-calling wiring test (#6836)
* studio: deterministic backend tool-calling wiring test Add a deterministic, download-free test that exercises the shared tool-calling seam both inference backends use. InferenceBackend (transformers) and MLXInferenceBackend both render the prompt through apply_chat_template_for_generation(..., tools=...) and stream cumulative text into run_safetensors_tool_loop. The existing test_safetensors_tool_loop.py covers the parser and the loop state machine with fake generators but does not cover the backend's own tool-injection seam, so a regression that drops the tool schema before the tokenizer, or fails to feed a tool result back into generation, would slip through. The test drives that seam with fakes: a tokenizer that records the tools it is handed, a canned tool-call generation, and a stub executor. It asserts the full chain: tools reach the chat template, the loop parses the call, the tool is dispatched once with the parsed arguments, the result is fed back, generation re-enters, and the final answer streams after the tool result. It also guards that the raw tool-call markup never leaks to the client as content. The test imports no torch, unsloth, or mlx, so it runs in the portable Backend CI alongside the tool-call parser tests and stays sub-second. Follow-up to the parser test PRs #5620 and #5704. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: assert the tool result is fed back before the final turn Strengthen the wiring test so single_turn records each turn's conversation and the test asserts the tool result message is present in the conversation handed to the final generation turn. Event ordering alone did not catch a loop that stops appending the tool output before re-entering generation, because the fake generation ignores the conversation; this closes that gap. * studio: tighten comments in tool-calling wiring test * studio: shorten comments in tool-calling wiring test --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
f38672da65
commit
e9f49c62dd
1 changed files with 179 additions and 0 deletions
179
studio/backend/tests/test_safetensors_toolcall_wiring.py
Normal file
179
studio/backend/tests/test_safetensors_toolcall_wiring.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path.
|
||||
|
||||
The parser and the cumulative-text state machine are already covered exhaustively by
|
||||
``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the
|
||||
*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and
|
||||
``MLXInferenceBackend`` render the prompt through the shared
|
||||
``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the
|
||||
shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and
|
||||
``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU
|
||||
test of that seam covers the macOS MLX path too).
|
||||
|
||||
This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the
|
||||
``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full
|
||||
agentic chain end to end:
|
||||
|
||||
tools injected into the template -> loop parses the call -> tool dispatched once ->
|
||||
tool result fed back -> generation re-entered -> final answer streamed.
|
||||
|
||||
It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling
|
||||
end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the
|
||||
tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704).
|
||||
"""
|
||||
|
||||
from core.inference.chat_template_helpers import apply_chat_template_for_generation
|
||||
from core.inference.safetensors_agentic import run_safetensors_tool_loop
|
||||
|
||||
TOOL_NAME = "get_weather"
|
||||
TOOL_ARGS = {"city": "Paris"}
|
||||
FAKE_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": TOOL_NAME,
|
||||
"description": "Get the current weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
# Full parser matrix lives in test_safetensors_tool_loop.py.
|
||||
TOOL_CALL_TEXT = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>'
|
||||
FINAL_ANSWER = "The weather in Paris is sunny and 22C."
|
||||
TOOL_RESULT = "Paris: sunny, 22C"
|
||||
|
||||
|
||||
class RecordingTokenizer:
|
||||
"""Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``.
|
||||
|
||||
Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the
|
||||
real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually
|
||||
forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen``
|
||||
holding ``None``.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.tools_seen: list = []
|
||||
self.call_count = 0
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.call_count += 1
|
||||
self.tools_seen.append(kwargs.get("tools"))
|
||||
return "PROMPT"
|
||||
|
||||
|
||||
class StubExecutor:
|
||||
"""Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result.
|
||||
|
||||
A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run.
|
||||
"""
|
||||
|
||||
def __init__(self, result: str):
|
||||
self.result = result
|
||||
self.calls: list[tuple[str, dict]] = []
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
name,
|
||||
arguments,
|
||||
*,
|
||||
cancel_event = None,
|
||||
timeout = None,
|
||||
session_id = None,
|
||||
rag_scope = None,
|
||||
disable_sandbox = False,
|
||||
):
|
||||
self.calls.append((name, arguments))
|
||||
return self.result
|
||||
|
||||
|
||||
def _collect(generator, max_events = 200):
|
||||
events = []
|
||||
for ev in generator:
|
||||
events.append(ev)
|
||||
if len(events) >= max_events:
|
||||
break
|
||||
return events
|
||||
|
||||
|
||||
def _tool_names(tools):
|
||||
return [(t.get("function") or {}).get("name") for t in (tools or [])]
|
||||
|
||||
|
||||
def test_backend_seam_injects_tools_and_drives_full_tool_loop():
|
||||
"""The shared backend seam forwards tools into the chat template, and the loop parses the call,
|
||||
dispatches it once, feeds the result back, and re-enters generation for the final answer."""
|
||||
tok = RecordingTokenizer()
|
||||
executor = StubExecutor(TOOL_RESULT)
|
||||
turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER])
|
||||
active_tools_seen: list = []
|
||||
conversations_seen: list = []
|
||||
|
||||
def single_turn(conversation, *, active_tools = None):
|
||||
# Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots.
|
||||
active_tools_seen.append(active_tools)
|
||||
conversations_seen.append([dict(m) for m in conversation])
|
||||
apply_chat_template_for_generation(tok, conversation, tools = active_tools)
|
||||
text = next(turns)
|
||||
mid = len(text) // 2
|
||||
acc = ""
|
||||
for chunk in (text[:mid], text[mid:]):
|
||||
acc += chunk
|
||||
yield acc
|
||||
|
||||
events = _collect(
|
||||
run_safetensors_tool_loop(
|
||||
single_turn = single_turn,
|
||||
messages = [{"role": "user", "content": "What is the weather in Paris?"}],
|
||||
tools = [FAKE_TOOL],
|
||||
execute_tool = executor,
|
||||
max_tool_iterations = 3,
|
||||
)
|
||||
)
|
||||
|
||||
# 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools).
|
||||
assert tok.tools_seen, "tokenizer.apply_chat_template was never called"
|
||||
assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer"
|
||||
assert TOOL_NAME in _tool_names(tok.tools_seen[0])
|
||||
|
||||
# 2. Loop offered the tool to the first generation turn.
|
||||
assert active_tools_seen and active_tools_seen[0] is not None
|
||||
assert TOOL_NAME in _tool_names(active_tools_seen[0])
|
||||
|
||||
# 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result.
|
||||
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
||||
tool_ends = [e for e in events if e["type"] == "tool_end"]
|
||||
assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME
|
||||
assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls
|
||||
assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT
|
||||
|
||||
# 6. Final answer streams after the tool result: loop appended it and re-entered generation.
|
||||
contents = [e for e in events if e["type"] == "content"]
|
||||
assert contents and FINAL_ANSWER in contents[-1]["text"]
|
||||
last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end")
|
||||
last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content")
|
||||
assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result"
|
||||
|
||||
# 6b. Tool result fed back into the conversation before the final turn (6 alone misses this:
|
||||
# the fake generation ignores the conversation).
|
||||
assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call"
|
||||
final_turn_convo = conversations_seen[1]
|
||||
assert any(
|
||||
TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo
|
||||
), "tool result was not fed back into the conversation before the final generation turn"
|
||||
|
||||
# 7. Guard: raw tool-call markup never leaked to the client as content.
|
||||
for e in contents:
|
||||
assert "<tool_call>" not in e["text"]
|
||||
assert TOOL_NAME not in e["text"]
|
||||
Loading…
Add table
Add a link
Reference in a new issue