vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-17 21:03:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 31

studio: route bare tool-call fragments to the call that owns the index now (#8755)

Browse source
This commit is contained in:
Maheswar Kumar 2026-08-16 16:34:58 +12:00 • committed by GitHub
parent 84c1887fcf
commit 321c3ee019
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 411 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

68
studio/backend/core/inference/llama_cpp.py
View file

@ -31,6 +31,7 @@ import uuid
from dataclasses import dataclass, replace
from pathlib import Path
from typing import (
Any,
Callable,
Collection,
Generator,
@ -20230,6 +20231,9 @@ class LlamaCppBackend:
in_thinking = False
has_content_tokens = False
tool_calls_acc = {} # Structured delta.tool_calls fragments
# accumulator key each index maps to now, (index, call_id) after a fork
tool_calls_open_key: dict[Any, Any] = {}
tool_calls_synthetic_ids: set[Any] = set() # slots still on a placeholder id
has_structured_tc = False
_iter_usage = None
_iter_timings = None
@ -20332,31 +20336,53 @@ class LlamaCppBackend:
yield {"type": "content", "text": cumulative_display}
for tc_d in tc_deltas:
idx = tc_d.get("index", 0)
if idx not in tool_calls_acc:
tool_calls_acc[idx] = {
"id": tc_d.get("id", f"call_{idx}"),
tc_id = tc_d.get("id")
# a second call at a reused index forks its own slot; bare fragments follow whichever call owns the index now
acc_key = tool_calls_open_key.get(idx, idx)
if isinstance(tc_id, str) and tc_id:
open_id = tool_calls_acc.get(acc_key, {}).get("id")
# a synthetic placeholder is not a rival call, so a first real id still lands on the slot that is open
if (
open_id
and acc_key not in tool_calls_synthetic_ids
and open_id != tc_id
):
first_id = tool_calls_acc.get(idx, {}).get("id")
acc_key = idx if first_id == tc_id else (idx, tc_id)
tool_calls_open_key[idx] = acc_key
if acc_key not in tool_calls_acc:
tool_calls_acc[acc_key] = {
"id": tc_id or f"call_{idx}",
"type": "function",
"function": {
"name": "",
"arguments": "",
},
}
elif tc_d.get("id"):
if not tc_id:
tool_calls_synthetic_ids.add(acc_key)
elif tc_id:
# Update ID if a real one
# arrives on a later delta.
tool_calls_acc[idx]["id"] = tc_d["id"]
tool_calls_acc[acc_key]["id"] = tc_id
tool_calls_synthetic_ids.discard(acc_key)
func = tc_d.get("function", {})
if func.get("name"):
tool_calls_acc[idx]["function"]["name"] += func["name"]
if func.get("arguments"):
tool_calls_acc[idx]["function"]["arguments"] += func[
"arguments"
tool_calls_acc[acc_key]["function"]["name"] += func[
"name"
]
current_name = tool_calls_acc[idx]["function"].get(
if func.get("arguments"):
tool_calls_acc[acc_key]["function"]["arguments"] += (
func["arguments"]
)
current_name = tool_calls_acc[acc_key]["function"].get(
"name", ""
)
fallback_id = f"call_{idx}"
current_id = tool_calls_acc[idx].get("id", fallback_id)
current_id = tool_calls_acc[acc_key].get("id", fallback_id)
is_first_accumulated_call = (
next(iter(tool_calls_acc)) == acc_key
)
already_started = (
current_id in provisional_started_tool_calls
)
@ -20390,7 +20416,7 @@ class LlamaCppBackend:
)
# Keep small-argument tools on the normal path.
_args_len = len(
tool_calls_acc[idx]["function"].get("arguments", "")
tool_calls_acc[acc_key]["function"].get("arguments", "")
)
_payload_is_large = (
current_name == "render_html"
@ -20398,7 +20424,10 @@ class LlamaCppBackend:
)
if (
current_name
and (idx == 0 or not disable_parallel_tool_use)
and (
is_first_accumulated_call
or not disable_parallel_tool_use
)
and has_real_id
and not already_started
and not _is_completed_one_shot
@ -20429,9 +20458,9 @@ class LlamaCppBackend:
if current_id in provisional_started_tool_calls:
if current_id not in arg_streamed_tool_call_ids:
arg_streamed_tool_call_ids.add(current_id)
_args_backlog = tool_calls_acc[idx]["function"].get(
"arguments", ""
)
_args_backlog = tool_calls_acc[acc_key][
"function"
].get("arguments", "")
if _args_backlog:
yield {
"type": "tool_args",
@ -20940,10 +20969,11 @@ class LlamaCppBackend:
if has_structured_tc:
# Drop incomplete fragments (e.g. from max_tokens
# truncation or disconnect).
# first-seen order: a call forked onto a reused index arrived after every call already open, so it must execute and replay last
tool_calls = [
tool_calls_acc[i]
for i in sorted(tool_calls_acc)
if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
call
for call in tool_calls_acc.values()
if (call.get("function", {}).get("name", "").strip())
] or None
if not tool_calls:
# Unconditional re-parse: we only reach DRAINING when the buffer looked like a

17
studio/backend/core/inference/studio_tool_loop.py
View file

@ -352,6 +352,9 @@ class _Turn:
by_index: dict[Any, dict[str, Any]] = field(default_factory = dict)
order: list[Any] = field(default_factory = list)
# call key each delta index maps to: the index itself until a second call
# forks off it, then (index, call_id).
open_key_by_index: dict[int, Any] = field(default_factory = dict)
last_index: int | None = None
round: int = 0
healed: list[dict[str, Any]] = field(default_factory = list)
@ -466,15 +469,23 @@ class _Turn:
# continue the call that is already open instead.
index = self.last_index if self.last_index is not None else len(self.order)
call_id = raw_call.get("id")
key: Any = index
# continue whichever call owns this index now: index restarts at 0
# for every tool round, so after a fork the bare argument fragments
# belong to the newer call.
key: Any = self.open_key_by_index.get(index, index)
if isinstance(call_id, str) and call_id:
open_id = self.by_index.get(index, {}).get("id")
open_id = self.by_index.get(key, {}).get("id")
if open_id and open_id != call_id:
# Two distinct calls reported at the same index. Merging them
# concatenates their argument JSON into one unparseable blob
# and loses an intent, so key the second on its own id.
key = (index, call_id)
# A fragment that names the call this index opened first goes
# back to it: an id beats the latest-index mapping, which only
# exists to place the fragments that carry no id.
first_id = self.by_index.get(index, {}).get("id")
key = index if first_id == call_id else (index, call_id)
self.last_index = index
self.open_key_by_index[index] = key
if key not in self.by_index:
self.by_index[key] = {
"id": "",

280
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -4939,3 +4939,283 @@ def test_provisional_text_card_closed_when_parse_fails(monkeypatch):
assert executed == [] # nothing parsed, nothing ran
assert ends, "provisional card left dangling (no tool_end)"
assert ends[-1]["tool_call_id"] == "call_0"
def _tool_call_fragment(
index: int,
arguments: str,
call_id: str | None = None,
) -> str:
delta: dict = {"index": index, "function": {"arguments": arguments}}
if call_id is not None:
delta["id"] = call_id
return _sse({"tool_calls": [delta]})
def _tool_call_opening(index: int, call_id: str, name: str, arguments: str) -> str:
return _sse(
{
"tool_calls": [
{
"index": index,
"id": call_id,
"type": "function",
"function": {"name": name, "arguments": arguments},
}
]
}
)
def test_second_structured_call_at_one_index_keeps_its_own_fragments(monkeypatch):
"""Two tool rounds in one llama-server response, both streamed at index 0.
llama-server restarts ``delta.tool_calls[].index`` at 0 for every round
while giving each call its own id, and the continuation fragments carrying
the rest of the arguments arrive bare. Keying the accumulator on the index
alone appended round two's name and argument tail to round one, leaving a
single ``web_searchweb_search`` entry that matched no enabled tool, so
neither call ran.
"""
stream = [
_tool_call_opening(0, "call_a", "web_search", '{"query":'),
_tool_call_fragment(0, '"first"}'),
_tool_call_opening(0, "call_b", "web_search", '{"query":'),
_tool_call_fragment(0, '"second"}'),
_finish("tool_calls"),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append({"name": name, "arguments": arguments})
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search twice"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 2,
)
)
assert calls == [
{"name": "web_search", "arguments": {"query": "first"}},
{"name": "web_search", "arguments": {"query": "second"}},
]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
"call_a",
"call_b",
]
# The replayed conversation must list both calls with their own arguments.
asst = next(m for m in payloads[1]["messages"] if m.get("tool_calls"))
assert [tc["id"] for tc in asst["tool_calls"]] == ["call_a", "call_b"]
assert [tc["function"]["arguments"] for tc in asst["tool_calls"]] == [
'{"query":"first"}',
'{"query":"second"}',
]
def test_structured_fragment_naming_its_call_goes_back_to_that_call(monkeypatch):
"""Two calls at index 0 with the id repeated on every argument fragment.
The latest-index mapping only exists to place fragments that carry no id,
so a fragment naming the call the index opened first has to go back to it
rather than fork a third slot. Forking left that call with truncated JSON
and dropped the fragment for having no function name.
"""
stream = [
_tool_call_opening(0, "call_a", "web_search", '{"query":'),
_tool_call_opening(0, "call_b", "web_search", '{"query":'),
_tool_call_fragment(0, '"first"}', call_id = "call_a"),
_tool_call_fragment(0, '"second"}', call_id = "call_b"),
_finish("tool_calls"),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append({"name": name, "arguments": arguments})
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search twice"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 2,
)
)
assert calls == [
{"name": "web_search", "arguments": {"query": "first"}},
{"name": "web_search", "arguments": {"query": "second"}},
]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
"call_a",
"call_b",
]
def test_structured_call_id_arriving_after_the_opening_delta_updates_that_call(monkeypatch):
"""llama-server can open a call with no id and send the real one later.
The opening slot holds the synthetic ``call_0`` until then, and treating
that placeholder as a rival id forked a second nameless slot: the fork was
dropped for having no function name and the original call ran on truncated
arguments.
"""
stream = [
_sse(
{
"tool_calls": [
{
"index": 0,
"type": "function",
"function": {"name": "web_search", "arguments": '{"query":'},
}
]
}
),
_tool_call_fragment(0, '"late"}', call_id = "call_late"),
_finish("tool_calls"),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append({"name": name, "arguments": arguments})
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 2,
)
)
assert calls == [{"name": "web_search", "arguments": {"query": "late"}}]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == ["call_late"]
def test_structured_call_forked_onto_a_reused_index_executes_last(monkeypatch):
"""A first round at indices 0 and 1, then a second round back at index 0.
The fork belongs at the end: it arrived after both first-round calls, and
running it ahead of the index-1 call reorders side effects for stateful
tools. Grouping every fork next to the index it reused did exactly that.
"""
stream = [
_tool_call_opening(0, "call_a", "web_search", '{"query":"a"}'),
_tool_call_opening(1, "call_b", "web_search", '{"query":"b"}'),
_tool_call_opening(0, "call_c", "web_search", '{"query":"c"}'),
_finish("tool_calls"),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append(arguments)
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search three times"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 2,
)
)
assert calls == [{"query": "a"}, {"query": "b"}, {"query": "c"}]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
"call_a",
"call_b",
"call_c",
]
asst = next(m for m in payloads[1]["messages"] if m.get("tool_calls"))
assert [tc["id"] for tc in asst["tool_calls"]] == ["call_a", "call_b", "call_c"]
def test_parallel_disabled_suppresses_provisional_for_reused_index(monkeypatch):
"""A later call at reused index 0 is not the first accumulated call.
``parallel_tool_calls=false`` permits a provisional card only for the call
that can execute. Checking the raw index admitted every fork at index 0,
leaving a card for the truncated call that could only close empty.
"""
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
big_cmd = "echo start\n" + "\n".join(f"echo line {i}" for i in range(60))
python_args = json.dumps({"code": big_code})
terminal_args = json.dumps({"command": big_cmd})
stream = [
_tool_call_opening(0, "call_py", "python", python_args[:-1]),
_tool_call_fragment(0, python_args[-1]),
_tool_call_opening(0, "call_term", "terminal", terminal_args[:-1]),
_tool_call_fragment(0, terminal_args[-1]),
_finish("tool_calls"),
_done(),
]
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream, final_stream], payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "OK"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "do the first only"}],
tools = [
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "terminal"}},
],
max_tool_iterations = 1,
disable_parallel_tool_use = True,
permission_mode = "off",
)
)
provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")]
assert [e["tool_call_id"] for e in provisional] == ["call_py"]
assert calls == [("python", {"code": big_code})]
assert not [
e
for e in events
if e.get("tool_call_id") == "call_term"
and e.get("type") in {"tool_start", "tool_args", "tool_end"}
]

68
studio/backend/tests/test_studio_tool_loop.py
View file

@ -1036,3 +1036,71 @@ def test_a_skipped_duplicate_closes_the_card_the_provider_already_painted(execut
# every tool_end has a matching tool_start.
starts = _events(lines, "tool_start")
assert [start["tool_call_id"] for start in starts] == ["call_a", "call_a"]
def test_a_second_call_at_one_index_keeps_its_own_argument_fragments(executed):
"""Two tool rounds in one response, both streamed at index 0.
Providers restart ``delta.tool_calls[].index`` at 0 for every round while
giving each call its own id, and the continuation fragments carrying the
rest of the arguments are sent bare. Routing those by index alone appended
round two's tail to round one, producing an unparseable blob and running
both tools on the wrong arguments.
"""
transport = FakeTransport(
[
[
_sse({"tool_calls": [_call_delta(0, "call_a", "web_search", '{"query":')]}),
_sse({"tool_calls": [{"index": 0, "function": {"arguments": '"first"}'}}]}),
_sse({"tool_calls": [_call_delta(0, "call_b", "web_search", '{"query":')]}),
_sse({"tool_calls": [{"index": 0, "function": {"arguments": '"second"}'}}]}),
_sse(finish = "tool_calls"),
_DONE,
],
[_sse({"content": "done"}), _sse(finish = "stop"), _DONE],
],
heals = False,
)
_run(transport)
assert [call["arguments"] for call in executed] == [{"query": "first"}, {"query": "second"}]
def test_a_fragment_naming_its_call_goes_back_to_that_call(executed):
"""Two calls at index 0 with the id repeated on every argument fragment.
The latest-index mapping only exists to place fragments that carry no id, so
a fragment that names the call the index opened first has to go back to it
rather than fork a third slot. Forking left that call with truncated JSON,
which reaches the tool as ``_raw``, and dropped the fragment for having no
function name.
"""
transport = FakeTransport(
[
[
_sse({"tool_calls": [_call_delta(0, "call_a", "web_search", '{"query":')]}),
_sse({"tool_calls": [_call_delta(0, "call_b", "web_search", '{"query":')]}),
_sse(
{
"tool_calls": [
{"index": 0, "id": "call_a", "function": {"arguments": '"first"}'}}
]
}
),
_sse(
{
"tool_calls": [
{"index": 0, "id": "call_b", "function": {"arguments": '"second"}'}}
]
}
),
_sse(finish = "tool_calls"),
_DONE,
],
[_sse({"content": "done"}), _sse(finish = "stop"), _DONE],
],
heals = False,
)
_run(transport)
assert [call["arguments"] for call in executed] == [{"query": "first"}, {"query": "second"}]