250 lines
8.9 KiB
Python
250 lines
8.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for the F3 validation-retry pass in ``run_safetensors_tool_loop``.
|
|
|
|
The pass detects two failure modes between parse and dispatch:
|
|
|
|
* The model emitted a tool call to a name not in the request's tools.
|
|
* The model emitted arguments that are not a JSON object (auto-heal off).
|
|
|
|
Either failure triggers a corrective tool-result message tied to the
|
|
hallucinated ``tool_call_id`` and re-enters the model loop. The pass is
|
|
bounded by ``max_validation_retries`` so the loop cannot retry forever.
|
|
"""
|
|
|
|
from core.inference.safetensors_agentic import run_safetensors_tool_loop
|
|
|
|
|
|
class FakeExecuteTool:
|
|
def __init__(self, results = None):
|
|
self.results = list(results or [])
|
|
self.calls = []
|
|
|
|
def __call__(
|
|
self, name, arguments, *, cancel_event = None, timeout = None, session_id = None
|
|
):
|
|
self.calls.append((name, arguments))
|
|
return self.results.pop(0) if self.results else "OK"
|
|
|
|
|
|
def _multi_turn(turns):
|
|
"""Yield cumulative-text generators, one per turn."""
|
|
turn_iter = iter(turns)
|
|
|
|
def _gen(_messages):
|
|
try:
|
|
chunks = next(turn_iter)
|
|
except StopIteration:
|
|
return
|
|
acc = ""
|
|
for c in chunks:
|
|
acc += c
|
|
yield acc
|
|
|
|
return _gen
|
|
|
|
|
|
def _collect(loop, cap = 200):
|
|
out = []
|
|
for ev in loop:
|
|
out.append(ev)
|
|
if len(out) >= cap:
|
|
break
|
|
return out
|
|
|
|
|
|
REAL_TOOLS = [
|
|
{"type": "function", "function": {"name": "web_search"}},
|
|
{"type": "function", "function": {"name": "python"}},
|
|
]
|
|
|
|
|
|
class TestUnknownTool:
|
|
def test_unknown_tool_triggers_retry(self):
|
|
# Turn 1: hallucinated tool name "missing_tool".
|
|
# Turn 2: valid call after the corrective nudge.
|
|
single_turn = _multi_turn(
|
|
[
|
|
['<tool_call>{"name":"missing_tool","arguments":{}}</tool_call>'],
|
|
[
|
|
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
|
|
],
|
|
["Done."],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool(["result"])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
)
|
|
)
|
|
# The bad call was NOT executed; the good follow-up call ran.
|
|
assert exec_fn.calls == [("web_search", {"query": "x"})]
|
|
# Final content carries the model's answer.
|
|
contents = [e for e in events if e["type"] == "content"]
|
|
assert any("Done" in e.get("text", "") for e in contents)
|
|
|
|
def test_retry_budget_exhausted_falls_through(self):
|
|
# Three unknown calls in a row. With max_validation_retries=2
|
|
# the first two are retried; the third falls through to the
|
|
# existing per-tool error path (which never executes a real
|
|
# tool but emits an error result to the model).
|
|
single_turn = _multi_turn(
|
|
[
|
|
['<tool_call>{"name":"missing_a","arguments":{}}</tool_call>'],
|
|
['<tool_call>{"name":"missing_b","arguments":{}}</tool_call>'],
|
|
['<tool_call>{"name":"missing_c","arguments":{}}</tool_call>'],
|
|
["Sorry, I cannot proceed."],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool([])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
max_validation_retries = 2,
|
|
)
|
|
)
|
|
# execute_tool is never called for an unknown name in either
|
|
# the F3 retry arm or the existing per-tool error arm.
|
|
assert exec_fn.calls == []
|
|
# The loop ultimately exits cleanly; we expect at least one
|
|
# tool_end event from the existing error path on the third turn.
|
|
kinds = [e.get("type") for e in events]
|
|
assert "status" in kinds
|
|
|
|
def test_retries_disabled_means_no_retry(self):
|
|
# With max_validation_retries=0, the F3 arm never engages and
|
|
# behavior matches the pre-F3 path: the existing per-tool
|
|
# error message is emitted but no corrective re-entry happens.
|
|
single_turn = _multi_turn(
|
|
[
|
|
['<tool_call>{"name":"missing","arguments":{}}</tool_call>'],
|
|
["bye"],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool([])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
max_validation_retries = 0,
|
|
)
|
|
)
|
|
# No retry round; the unknown call goes straight to the
|
|
# existing error path.
|
|
assert exec_fn.calls == []
|
|
tool_ends = [e for e in events if e["type"] == "tool_end"]
|
|
# Existing per-tool error path emits a tool_end with an Error
|
|
# result so the model sees the failure.
|
|
assert any("not enabled" in str(e.get("result", "")) for e in tool_ends)
|
|
|
|
|
|
class TestMalformedArgs:
|
|
def test_malformed_args_bypassed_when_heal_on(self):
|
|
# With auto_heal_tool_calls=True (the default), string args are
|
|
# healed to {"query": "..."} for web_search. F3 leaves heal
|
|
# behavior intact and only catches strictly-impossible shapes.
|
|
single_turn = _multi_turn(
|
|
[
|
|
[
|
|
'<tool_call>{"name":"web_search","arguments":"some text"}</tool_call>'
|
|
],
|
|
["all done"],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool(["ok"])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
auto_heal_tool_calls = True,
|
|
)
|
|
)
|
|
# The heal path coerces "some text" to {"query": "some text"}
|
|
# and execute_tool runs.
|
|
assert len(exec_fn.calls) == 1
|
|
assert exec_fn.calls[0][0] == "web_search"
|
|
|
|
def test_malformed_args_caught_when_heal_off(self):
|
|
# With auto_heal off, a non-dict arguments value is a hard
|
|
# malformed-args failure and the F3 arm catches it.
|
|
single_turn = _multi_turn(
|
|
[
|
|
[
|
|
'<tool_call>{"name":"web_search","arguments":"not a dict"}</tool_call>'
|
|
],
|
|
[
|
|
'<tool_call>{"name":"web_search","arguments":{"query":"sf"}}</tool_call>'
|
|
],
|
|
["sunny"],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool(["sunny in sf"])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
auto_heal_tool_calls = False,
|
|
)
|
|
)
|
|
# The first call was caught by F3, not executed. The second
|
|
# (well-formed) call ran.
|
|
assert exec_fn.calls == [("web_search", {"query": "sf"})]
|
|
|
|
|
|
class TestNoOpCases:
|
|
def test_no_tools_no_validation(self):
|
|
# Empty tools list means allowed_tool_names is empty, so the
|
|
# F3 arm never engages.
|
|
single_turn = _multi_turn(
|
|
[
|
|
['<tool_call>{"name":"anything","arguments":{}}</tool_call>'],
|
|
["done"],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool([])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = [],
|
|
execute_tool = exec_fn,
|
|
)
|
|
)
|
|
# No allowlist gate, no F3 retry; the bad call goes through
|
|
# the dispatch loop and execute_tool runs.
|
|
assert exec_fn.calls == [("anything", {})]
|
|
|
|
def test_valid_call_no_retry_overhead(self):
|
|
# A clean valid call does not trigger F3 at all.
|
|
single_turn = _multi_turn(
|
|
[
|
|
[
|
|
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
|
|
],
|
|
["final"],
|
|
]
|
|
)
|
|
exec_fn = FakeExecuteTool(["res"])
|
|
events = _collect(
|
|
run_safetensors_tool_loop(
|
|
single_turn = single_turn,
|
|
messages = [{"role": "user", "content": "hi"}],
|
|
tools = REAL_TOOLS,
|
|
execute_tool = exec_fn,
|
|
)
|
|
)
|
|
assert exec_fn.calls == [("web_search", {"query": "x"})]
|