ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Gemma 4 GGUF OpenAI API streams

Browse source
This commit is contained in:
wasimysaid 2026-06-19 16:27:33 +02:00
commit 165f4838e3
10 changed files with 803 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

168
studio/backend/core/inference/tool_call_parser.py
View file

@ -16,16 +16,18 @@ import re
# so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<\|tool_call>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
# Prefixes the streaming buffer watches for to gate in-progress text.
TOOL_XML_SIGNALS = ("<tool_call>", "<function=")
TOOL_XML_SIGNALS = ("<tool_call>", "<|tool_call>", "<function=")
# Nudges + error prefixes shared by the GGUF and safetensors loops.
@ -76,14 +78,17 @@ RAG_SEARCH_CAP_NUDGE = (
# Pre-compiled patterns reused by ``parse_tool_calls_from_text``.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
# [\w-] so hyphenated MCP param names (issue-number) aren't dropped.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_PARAM_CLOSE_TAG = "</parameter>"
_FUNC_CLOSE_TAG = "</function>"
_GEMMA_QUOTE = '<|"|>'
def _inside_open_parameter(content: str, pos: int) -> bool:
@ -111,6 +116,116 @@ def strip_tool_markup(text: str, *, final: bool = False) -> str:
return text.strip() if final else text
def _balanced_brace_end(content: str, brace_start: int) -> int:
depth = 0
i = brace_start
in_string = False
in_gemma_string = False
while i < len(content):
if content.startswith(_GEMMA_QUOTE, i):
in_gemma_string = not in_gemma_string
i += len(_GEMMA_QUOTE)
continue
ch = content[i]
if in_gemma_string:
i += 1
continue
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return i
i += 1
return -1
def _normalise_gemma_quoted_strings(src: str) -> str:
parts: list[str] = []
i = 0
while i < len(src):
if not src.startswith(_GEMMA_QUOTE, i):
parts.append(src[i])
i += 1
continue
end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE))
if end < 0:
parts.append(src[i:])
break
raw_value = src[i + len(_GEMMA_QUOTE) : end]
parts.append(json.dumps(raw_value))
i = end + len(_GEMMA_QUOTE)
return "".join(parts)
def _quote_gemma_object_keys(src: str) -> str:
parts: list[str] = []
i = 0
in_string = False
while i < len(src):
ch = src[i]
if in_string:
parts.append(ch)
if ch == "\\" and i + 1 < len(src):
parts.append(src[i + 1])
i += 2
continue
if ch == '"':
in_string = False
i += 1
continue
if ch == '"':
in_string = True
parts.append(ch)
i += 1
continue
if ch not in "{,":
parts.append(ch)
i += 1
continue
parts.append(ch)
i += 1
key_start = i
while i < len(src) and src[i].isspace():
i += 1
key_name_start = i
while i < len(src) and (src[i].isalnum() or src[i] in "_-"):
i += 1
key_name = src[key_name_start:i]
colon_pos = i
while colon_pos < len(src) and src[colon_pos].isspace():
colon_pos += 1
if key_name and colon_pos < len(src) and src[colon_pos] == ":":
parts.append(src[key_start:key_name_start])
parts.append(json.dumps(key_name))
parts.append(src[i:colon_pos])
parts.append(":")
i = colon_pos + 1
else:
parts.append(src[key_start:i])
return "".join(parts)
def _gemma_arguments_to_json(args_src: str) -> dict:
"""Parse Gemma 4's native call:name{key:value} argument object."""
args_src = args_src.strip()
if not args_src:
return {}
src = _normalise_gemma_quoted_strings(args_src)
src = "{" + src + "}"
src = _quote_gemma_object_keys(src)
return json.loads(src)
def parse_tool_calls_from_text(
content: str,
*,
@ -123,10 +238,12 @@ def parse_tool_calls_from_text(
dicts. ``arguments`` is always a JSON string so callers can hand it
straight back into an OpenAI-style response.
Handles two shapes:
Handles three shapes:
- JSON inside ``<tool_call>`` tags:
``<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>``
- Gemma 4 native call blocks:
``<|tool_call>call:web_search{query:"..." }<tool_call|>``
- XML-style function blocks:
``<function=name><parameter=k>v</parameter></function>``
@ -141,26 +258,8 @@ def parse_tool_calls_from_text(
# JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth != 0:
i = _balanced_brace_end(content, brace_start)
if i < 0:
continue
if not allow_incomplete:
tail_after_json = content[i + 1 :].lstrip()
@ -183,6 +282,31 @@ def parse_tool_calls_from_text(
except (json.JSONDecodeError, ValueError):
pass
# Pattern 1b: Gemma 4 native call block:
# <|tool_call>call:terminal{command:"ls"}<tool_call|>
for m in _TC_GEMMA_START_RE.finditer(content):
brace_start = m.end() - 1
i = _balanced_brace_end(content, brace_start)
if i < 0:
continue
if not allow_incomplete:
tail_after_json = content[i + 1 :].lstrip()
if _TC_GEMMA_END_TAG_RE.match(tail_after_json) is None:
continue
try:
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": m.group(1),
"arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])),
},
}
)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: <function=name><parameter=k>v... -- closing tags optional;
# </function> isn't a body boundary since code values can contain it.
if not tool_calls:

156
studio/backend/core/tool_healing.py
View file

@ -19,20 +19,133 @@ import re
# issue-number) parse alongside the built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<\|tool_call>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_GEMMA_QUOTE = '<|"|>'
def _balanced_brace_end(content: str, brace_start: int) -> int:
depth = 0
i = brace_start
in_string = False
in_gemma_string = False
while i < len(content):
if content.startswith(_GEMMA_QUOTE, i):
in_gemma_string = not in_gemma_string
i += len(_GEMMA_QUOTE)
continue
ch = content[i]
if in_gemma_string:
i += 1
continue
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return i
i += 1
return -1
def _normalise_gemma_quoted_strings(src: str) -> str:
parts: list[str] = []
i = 0
while i < len(src):
if not src.startswith(_GEMMA_QUOTE, i):
parts.append(src[i])
i += 1
continue
end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE))
if end < 0:
parts.append(src[i:])
break
raw_value = src[i + len(_GEMMA_QUOTE) : end]
parts.append(json.dumps(raw_value))
i = end + len(_GEMMA_QUOTE)
return "".join(parts)
def _quote_gemma_object_keys(src: str) -> str:
parts: list[str] = []
i = 0
in_string = False
while i < len(src):
ch = src[i]
if in_string:
parts.append(ch)
if ch == "\\" and i + 1 < len(src):
parts.append(src[i + 1])
i += 2
continue
if ch == '"':
in_string = False
i += 1
continue
if ch == '"':
in_string = True
parts.append(ch)
i += 1
continue
if ch not in "{,":
parts.append(ch)
i += 1
continue
parts.append(ch)
i += 1
key_start = i
while i < len(src) and src[i].isspace():
i += 1
key_name_start = i
while i < len(src) and (src[i].isalnum() or src[i] in "_-"):
i += 1
key_name = src[key_name_start:i]
colon_pos = i
while colon_pos < len(src) and src[colon_pos].isspace():
colon_pos += 1
if key_name and colon_pos < len(src) and src[colon_pos] == ":":
parts.append(src[key_start:key_name_start])
parts.append(json.dumps(key_name))
parts.append(src[i:colon_pos])
parts.append(":")
i = colon_pos + 1
else:
parts.append(src[key_start:i])
return "".join(parts)
def _gemma_arguments_to_json(args_src: str) -> dict:
args_src = args_src.strip()
if not args_src:
return {}
src = _normalise_gemma_quoted_strings(args_src)
src = "{" + src + "}"
src = _quote_gemma_object_keys(src)
return json.loads(src)
def parse_tool_calls_from_text(content: str) -> list[dict]:
@ -41,6 +154,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
Handles formats like:
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<|tool_call>call:web_search{query:"..."}<tool_call|>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
Closing tags (</tool_call>, </function>, </parameter>) are all
optional since models frequently omit them.
@ -51,26 +165,8 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
# skips braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
ch = content[i]
if in_string:
if ch == "\\" and i + 1 < len(content):
i += 2 # skip escaped character
continue
if ch == '"':
in_string = False
elif ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
break
i += 1
if depth == 0:
i = _balanced_brace_end(content, brace_start)
if i >= 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
@ -88,6 +184,26 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
except (json.JSONDecodeError, ValueError):
pass
# Pattern 1b: Gemma 4 native <|tool_call>call:name{key:value}<tool_call|>.
for m in _TC_GEMMA_START_RE.finditer(content):
brace_start = m.end() - 1
i = _balanced_brace_end(content, brace_start)
if i < 0:
continue
try:
tool_calls.append(
{
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {
"name": m.group(1),
"arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])),
},
}
)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: XML-style <function=name><parameter=key>value</parameter></function>
# All closing tags optional; models frequently omit them.
if not tool_calls:

4
studio/backend/models/inference.py
View file

@ -1101,6 +1101,8 @@ class ChoiceDelta(BaseModel):
role: Optional[str] = None
content: Optional[str] = None
reasoning_content: Optional[str] = None
tool_calls: Optional[list[dict]] = None
OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
@ -1136,6 +1138,8 @@ class CompletionMessage(BaseModel):
role: Literal["assistant"] = "assistant"
content: str
refusal: Optional[str] = None
reasoning_content: Optional[str] = None
tool_calls: Optional[list[dict]] = None
class CompletionChoice(BaseModel):

109
studio/backend/routes/inference.py
View file

@ -787,7 +787,10 @@ async def _aiter_llama_stream_items(
raise httpx.ReadTimeout("The model did not produce a first token in time.")
if response is not None:
_set_stream_response_read_timeout(response, remaining_s)
item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s)
# Keep httpx/httpcore's AnyIO cancel scope in this task.
# asyncio.wait_for would drive __anext__ in a child task.
async with asyncio.timeout(remaining_s):
item = await async_iter.__anext__()
else:
item = await async_iter.__anext__()
except asyncio.TimeoutError as exc:
@ -1286,6 +1289,7 @@ _TOOL_XML_RE = _re.compile(
# Hyphen in the name char-class matches MCP tool names with dashes
# (mcp__srv__list-issues) that would otherwise leak past this strip.
r"<(?:tool_call|function=[\w-]+)>.*?(?:</(?:tool_call|function)>|\Z)"
r"|<\|tool_call>.*?(?:<tool_call\|>|\Z)"
r"|</(?:tool_call|function)>"
r"|</parameter>\s*\Z",
_re.DOTALL,
@ -4842,6 +4846,28 @@ async def openai_chat_completions(
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
def _new_chat_reasoning_extractor():
return _ResponsesReasoningExtractor(
parse_think_markers = _responses_should_parse_think_markers(
payload,
llama_backend,
)
)
def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str:
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = delta,
finish_reason = finish_reason,
)
],
)
return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
# ── Tool-calling path (agentic loop) ──────────────────
# `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools`
# hard-override the per-request value, else falls back to
@ -5002,6 +5028,7 @@ async def openai_chat_completions(
# stays free for disconnect detection.
gen = gguf_generate_with_tools()
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
_stream_usage = None
_stream_timings = None
_stream_finish = None
@ -5024,6 +5051,7 @@ async def openai_chat_completions(
# streams cleanly.
if not event["text"]:
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
# Emit tool status as a custom SSE event (including
# empty ones to clear UI badges)
status_data = json.dumps(
@ -5038,6 +5066,7 @@ async def openai_chat_completions(
if event["type"] in ("tool_start", "tool_end"):
if event["type"] == "tool_start":
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
yield f"data: {json.dumps(event)}\n\n"
continue
@ -5059,19 +5088,23 @@ async def openai_chat_completions(
prev_text = clean_cumulative
if not new_text:
continue
api_monitor.append_reply(monitor_id, new_text)
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(content = new_text),
finish_reason = None,
)
],
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = reasoning_delta)
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta))
final_reasoning, final_visible = reasoning_extractor.finish()
if final_reasoning:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = final_reasoning)
)
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
if final_visible:
api_monitor.append_reply(monitor_id, final_visible)
yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible))
final_chunk = ChatCompletionChunk(
id = completion_id,
@ -5186,6 +5219,7 @@ async def openai_chat_completions(
# stays free for disconnect detection.
gen = gguf_generate()
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
_stream_usage = None
_stream_timings = None
_stream_finish = None
@ -5219,19 +5253,23 @@ async def openai_chat_completions(
prev_text = cumulative
if not new_text:
continue
api_monitor.append_reply(monitor_id, new_text)
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(content = new_text),
finish_reason = None,
)
],
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = reasoning_delta)
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta))
final_reasoning, final_visible = reasoning_extractor.finish()
if final_reasoning:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = final_reasoning)
)
yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
if final_visible:
api_monitor.append_reply(monitor_id, final_visible)
yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible))
# Final chunk
final_chunk = ChatCompletionChunk(
@ -5309,14 +5347,24 @@ async def openai_chat_completions(
continue
full_text = token
reasoning_text, visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _responses_should_parse_think_markers(
payload,
llama_backend,
),
)
message_kwargs = {"content": visible_text}
if reasoning_text:
message_kwargs["reasoning_content"] = reasoning_text
_choices.append(
CompletionChoice(
index = _idx,
message = CompletionMessage(content = full_text),
message = CompletionMessage(**message_kwargs),
finish_reason = _clamp_finish_reason(completion_finish),
)
)
_monitor_replies.append(full_text)
_monitor_replies.append(visible_text)
if completion_usage:
# The prompt is shared across all n choices, so count its
# tokens ONCE (OpenAI bills only generated tokens for each
@ -5338,7 +5386,7 @@ async def openai_chat_completions(
prompt_tokens_details = _prompt_tokens_details(_prompt_details),
),
)
monitor_reply = full_text
monitor_reply = _monitor_replies[-1] if _monitor_replies else ""
if _n > 1:
monitor_reply = "\n\n".join(
f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies)
@ -6686,8 +6734,9 @@ def _responses_should_parse_think_markers(
if llama_backend is not None and getattr(llama_backend, "is_loaded", False):
if getattr(llama_backend, "reasoning_always_on", False):
return True
if not getattr(llama_backend, "supports_reasoning", False):
return False
if getattr(llama_backend, "supports_reasoning", False):
return True
return False
if chat_req.enable_thinking is True:
return True
return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none")

29
studio/backend/tests/test_llama_route_timeouts.py
View file

@ -40,6 +40,35 @@ def test_stream_first_item_deadline_after_headers():
asyncio.run(_run())
def test_stream_first_item_deadline_does_not_hop_tasks():
async def _run():
outer_task = asyncio.current_task()
seen_tasks = []
class _One:
def __init__(self):
self.done = False
async def __anext__(self):
seen_tasks.append(asyncio.current_task())
if self.done:
raise StopAsyncIteration
self.done = True
return "data: {}"
out = []
async for item in inf_mod._aiter_llama_stream_items(
_One(),
first_token_deadline = time.monotonic() + 1,
):
out.append(item)
assert out == ["data: {}"]
assert seen_tasks == [outer_task, outer_task]
asyncio.run(_run())
def test_preheader_send_cleanup_on_disconnect_and_cancel():
async def _run(cancel_parent):
state = SimpleNamespace(disconnected = False, closed = False, cancelled = False)

23
studio/backend/tests/test_mcp_servers.py
View file

@ -423,6 +423,29 @@ def test_tool_healing_strip_handles_hyphenated_function_names():
assert out == "before after"
def test_tool_healing_strip_handles_gemma_native_tool_call():
from core.tool_healing import strip_tool_call_markup
out = strip_tool_call_markup(
'before <|tool_call>call:mcp__srv__list-issues{repo:"octocat/hello"}<tool_call|> after'
)
assert out == "before after"
def test_tool_healing_parser_handles_gemma_native_windows_path():
from core.tool_healing import parse_tool_calls_from_text
import json as _json
calls = parse_tool_calls_from_text(
r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}<tool_call|>'
)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "ls"
assert _json.loads(calls[0]["function"]["arguments"]) == {
"path": r"C:\Users\wasim\repo"
}
def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
"""A tool call not in the per-request list must be refused by the GGUF
agentic loop (mirroring the safetensors path)."""

261
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -1245,6 +1245,24 @@ class TestGgufVisionToolRouting:
return TestGgufVisionToolRouting._drive(_consume())
@staticmethod
def _sse_payloads(chunks):
payloads = []
for chunk in chunks:
if isinstance(chunk, bytes):
chunk = chunk.decode()
for line in str(chunk).splitlines():
if not line.startswith("data: "):
continue
data = line.removeprefix("data: ")
if data == "[DONE]":
continue
try:
payloads.append(json.loads(data))
except json.JSONDecodeError:
pass
return payloads
def test_image_request_with_enabled_tools_enters_gguf_tool_loop(self, monkeypatch):
import routes.inference as inf_mod
@ -1390,6 +1408,249 @@ class TestGgufVisionToolRouting:
assert "confirm_tool_calls requires stream=true" in entry["error"]
assert monitor.active_count() == 0
def test_standard_gguf_stream_splits_reasoning_content(self, monkeypatch):
import routes.inference as inf_mod
reset_tool_policy()
def _generate(**_kwargs):
yield "<thi"
yield "<think>plan"
yield "<think>plan</think>vis"
yield "<think>plan</think>visible"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"finish_reason": "stop",
}
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = False,
supports_reasoning = True,
reasoning_always_on = True,
_is_audio = False,
model_identifier = "test-gguf",
context_length = 4096,
generate_chat_completion = _generate,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
stream = True,
messages = [{"role": "user", "content": "hi"}],
)
response = self._drive(
openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
payloads = self._sse_payloads(self._consume_response(response))
deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")]
assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan"
assert "".join(d.get("content", "") for d in deltas) == "visible"
assert all("<think>" not in d.get("content", "") for d in deltas)
[entry] = monitor.snapshot()
assert entry["reply"] == "visible"
def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch):
import routes.inference as inf_mod
reset_tool_policy()
def _generate(**_kwargs):
yield "<think>plan</think>visible"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"finish_reason": "stop",
}
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = False,
supports_reasoning = True,
reasoning_always_on = False,
_is_audio = False,
model_identifier = "test-gguf",
context_length = 4096,
generate_chat_completion = _generate,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
stream = True,
messages = [{"role": "user", "content": "hi"}],
)
response = self._drive(
openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
payloads = self._sse_payloads(self._consume_response(response))
deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")]
assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan"
assert "".join(d.get("content", "") for d in deltas) == "visible"
[entry] = monitor.snapshot()
assert entry["reply"] == "visible"
def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled(
self, monkeypatch
):
import routes.inference as inf_mod
reset_tool_policy()
def _generate(**_kwargs):
yield "<think>leaked</think>visible"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"finish_reason": "stop",
}
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = False,
supports_reasoning = True,
reasoning_always_on = False,
_is_audio = False,
model_identifier = "test-gguf",
context_length = 4096,
generate_chat_completion = _generate,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
stream = True,
enable_thinking = False,
messages = [{"role": "user", "content": "hi"}],
)
response = self._drive(
openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
payloads = self._sse_payloads(self._consume_response(response))
deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")]
assert "".join(d.get("reasoning_content", "") for d in deltas) == "leaked"
assert "".join(d.get("content", "") for d in deltas) == "visible"
assert all("<think>" not in d.get("content", "") for d in deltas)
[entry] = monitor.snapshot()
assert entry["reply"] == "visible"
def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker(self, monkeypatch):
import routes.inference as inf_mod
reset_tool_policy()
def _plain(**_kwargs):
raise AssertionError("plain GGUF path should not be used")
def _tools(**_kwargs):
yield {
"type": "content",
"text": '<think>plan</think>visible <|tool_call>call:terminal{command:"ls"}<tool_call|>',
}
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"finish_reason": "stop",
}
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = True,
supports_reasoning = True,
reasoning_always_on = True,
_is_audio = False,
model_identifier = "test-gguf",
context_length = 4096,
generate_chat_completion = _plain,
generate_chat_completion_with_tools = _tools,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
stream = True,
enable_tools = True,
enabled_tools = ["terminal"],
messages = [{"role": "user", "content": "list files"}],
)
response = self._drive(
openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
payloads = self._sse_payloads(self._consume_response(response))
deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")]
assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan"
combined_content = "".join(d.get("content", "") for d in deltas)
assert combined_content == "visible "
assert "<|tool_call>" not in combined_content
[entry] = monitor.snapshot()
assert entry["reply"] == "visible "
def test_non_streaming_gguf_splits_reasoning_content(self, monkeypatch):
import routes.inference as inf_mod
reset_tool_policy()
def _generate(**_kwargs):
yield "<think>plan</think>visible"
yield {
"type": "metadata",
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"finish_reason": "stop",
}
backend = SimpleNamespace(
is_loaded = True,
is_vision = False,
supports_tools = False,
supports_reasoning = True,
reasoning_always_on = True,
_is_audio = False,
model_identifier = "test-gguf",
context_length = 4096,
generate_chat_completion = _generate,
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
payload = ChatCompletionRequest(
model = "default",
messages = [{"role": "user", "content": "hi"}],
)
response = self._drive(
openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
body = json.loads(response.body)
message = body["choices"][0]["message"]
assert message["content"] == "visible"
assert message["reasoning_content"] == "plan"
[entry] = monitor.snapshot()
assert entry["reply"] == "visible"
def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch):
import routes.inference as inf_mod

52
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -927,6 +927,38 @@ class TestResponsesNonStreamingAdapter:
assert [item["type"] for item in body["output"]] == ["message"]
assert body["output"][0]["content"][0]["text"] == "show <think>x</think> tags"
def test_reasoning_capable_gguf_parses_think_tags_by_default(self, monkeypatch):
body = self._run_with_message(
monkeypatch,
{"content": "<think>plan</think>answer"},
llama_backend = SimpleNamespace(
is_loaded = True,
reasoning_always_on = False,
supports_reasoning = True,
),
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}]
assert body["output"][1]["content"][0]["text"] == "answer"
def test_reasoning_capable_gguf_sanitizes_think_tags_when_disabled(self, monkeypatch):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"})
body = self._run_with_message(
monkeypatch,
{"content": "<think>leaked</think>answer"},
payload = payload,
llama_backend = SimpleNamespace(
is_loaded = True,
reasoning_always_on = False,
supports_reasoning = True,
),
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "leaked"}]
assert body["output"][1]["content"][0]["text"] == "answer"
def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch):
body = self._run_with_message(
monkeypatch,
@ -1332,10 +1364,10 @@ class TestResponsesStreamAdapter:
assert entry["status"] == "completed"
assert entry["reply"] == "plan"
def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch):
def test_reasoning_capable_gguf_stream_parses_think_tags_by_default(self, monkeypatch):
chunks = [
{"choices": [{"delta": {"content": "show <thi"}}]},
{"choices": [{"delta": {"content": "nk>x</think> tags"}}]},
{"choices": [{"delta": {"content": "<thi"}}]},
{"choices": [{"delta": {"content": "nk>plan</think>answer"}}]},
{"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
]
self._install_stream_mock(monkeypatch, chunks)
@ -1350,13 +1382,15 @@ class TestResponsesStreamAdapter:
reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta")
text_deltas = self._payloads(lines, "response.output_text.delta")
assert reasoning_deltas == []
assert "".join(event["delta"] for event in text_deltas) == "show <think>x</think> tags"
assert "".join(event["delta"] for event in reasoning_deltas) == "plan"
assert "".join(event["delta"] for event in text_deltas) == "answer"
completed = self._payloads(lines, "response.completed")[0]
assert [item["type"] for item in completed["response"]["output"]] == ["message"]
assert completed["response"]["output"][0]["content"][0]["text"] == (
"show <think>x</think> tags"
)
assert [item["type"] for item in completed["response"]["output"]] == [
"reasoning",
"message",
]
assert completed["response"]["output"][0]["content"][0]["text"] == "plan"
assert completed["response"]["output"][1]["content"][0]["text"] == "answer"
def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch):
chunks = [

45
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -10,6 +10,7 @@ calls, tool-result feedback, bad-JSON heal, duplicate-call short-circuit,
``__IMAGES__`` sentinel stripping, executor errors, cancel, and the iteration cap.
"""
import json
import threading
from typing import cast
@ -62,6 +63,46 @@ class TestParser:
assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
def test_gemma_native_tool_call(self):
text = '<|tool_call>call:terminal{command:"ls -la",workdir:"."}<tool_call|>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "terminal"
args = json.loads(result[0]["function"]["arguments"])
assert args == {"command": "ls -la", "workdir": "."}
def test_gemma_native_tool_call_template_quotes(self):
text = '<|tool_call>call:web_search{query:<|"|>openai news<|"|>}<tool_call|>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
assert json.loads(result[0]["function"]["arguments"]) == {"query": "openai news"}
def test_gemma_native_tool_call_template_quotes_escape_backslashes(self):
text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}<tool_call|>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "ls"
assert json.loads(result[0]["function"]["arguments"]) == {
"path": r"C:\Users\wasim\repo"
}
def test_gemma_native_tool_call_hyphenated_argument_name(self):
text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}<tool_call|>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "mcp__srv__create-issue"
assert json.loads(result[0]["function"]["arguments"]) == {"issue-title": "Bug report"}
def test_gemma_native_tool_call_keeps_braces_inside_string_value(self):
text = '<|tool_call>call:terminal{command:"echo {foo:bar}"}<tool_call|>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "terminal"
assert json.loads(result[0]["function"]["arguments"]) == {
"command": "echo {foo:bar}"
}
def test_xml_function_call(self):
text = "<function=python><parameter=code>print('hi')</parameter></function>"
result = parse_tool_calls_from_text(text)
@ -121,6 +162,7 @@ class TestParser:
def test_has_tool_signal(self):
assert has_tool_signal("blah <tool_call> x")
assert has_tool_signal("blah <|tool_call>call:terminal")
assert has_tool_signal("hi <function=foo>...")
assert not has_tool_signal("hello world")
@ -139,6 +181,8 @@ class TestParser:
def test_strip_markup_closed(self):
text = "before <tool_call>{}</tool_call> after"
assert strip_tool_markup(text) == "before after"
text = 'before <|tool_call>call:terminal{command:"ls"}<tool_call|> after'
assert strip_tool_markup(text) == "before after"
def test_strip_markup_unclosed_final(self):
text = "before <tool_call>{partial"
@ -146,6 +190,7 @@ class TestParser:
assert strip_tool_markup(text, final = True) == "before"
# Without final=True the unclosed run is preserved.
assert "partial" in strip_tool_markup(text)
assert strip_tool_markup("before <|tool_call>call:terminal{", final = True) == "before"
def test_streaming_strip_respects_disabled_healing(self):
raw = 'before <tool_call>{"name":"web_search"'

37
studio/backend/tests/test_tool_call_parser_strict.py
View file

@ -106,6 +106,43 @@ class TestParityWithJsonStyle:
assert json.loads(js[0]["function"]["arguments"]) == {"query": q}
class TestGemmaNativeStyle:
def test_closed_native_call_with_trailing_prose_is_accepted(self):
text = (
'<|tool_call>call:terminal{command:"ls -la",workdir:"."}<tool_call|>'
" running it now"
)
calls = parse_tool_calls_from_text(text, allow_incomplete = False)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "terminal"
assert json.loads(calls[0]["function"]["arguments"]) == {
"command": "ls -la",
"workdir": ".",
}
def test_unclosed_native_call_requires_healing(self):
text = '<|tool_call>call:terminal{command:"ls"}'
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "terminal"
def test_hyphenated_native_argument_name_is_accepted(self):
text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}<tool_call|>'
calls = parse_tool_calls_from_text(text, allow_incomplete = False)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "mcp__srv__create-issue"
assert json.loads(calls[0]["function"]["arguments"]) == {"issue-title": "Bug report"}
def test_native_template_quotes_preserve_windows_path(self):
text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}<tool_call|>'
calls = parse_tool_calls_from_text(text, allow_incomplete = False)
assert len(calls) == 1
assert json.loads(calls[0]["function"]["arguments"]) == {
"path": r"C:\Users\wasim\repo"
}
class TestHealingPathUnaffected:
def test_auto_heal_still_repairs_unclosed_function(self):
text = "<function=web_search><parameter=query>cats"