diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b4c908b0cb..2d30265abe 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,6 +564,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for n, tids in b["module_import_targets"].items(): if tids & after_used: continue # resolved -> fine + # `from __future__ import ...` is a compiler directive whose name is never loaded; skip it. + if all(t.startswith("from:__future__:") for t in tids): + continue newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -588,9 +591,19 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] # package object and only *add* submodule attributes (e.g. adding # `import urllib.error` next to `import urllib.request`). Nothing the name # resolved to before is lost, so no reference is re-pointed -- skip it. + # + # A deliberate *relocation* is also benign: a name's import source moves A -> B in + # THIS diff (old `from A import x` removed, new `from B import x` added). Mirrors the + # TARGET-MISSING tolerance. Re-pointing to a pre-existing target (clash) is NOT exempted. + removed_module_targets = before_module_targets - after_module_targets for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) if tbefore and tbefore != tafter and (tbefore - tafter): + lost = tbefore - tafter + gained = tafter - tbefore + relocated = lost <= removed_module_targets and gained <= added_module_targets + if relocated: + continue findings.append( ( "BLOCKER", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ccfc5cdfe..5e67f6b484 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -38,9 +38,21 @@ from core.inference.llama_server_args import ( strip_shadowing_flags, strip_split_mode_only, ) -from core.tool_healing import ( + +# Share strip / signal constants with the multi-format parser so BUFFERING also +# catches Llama-3 / Mistral / Gemma 4. +from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, - strip_tool_call_markup, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_mistral_closed_calls, + TOOL_XML_SIGNALS as _SHARED_TOOL_XML_SIGNALS, + RAG_MAX_SEARCHES_PER_TURN, + RAG_SEARCH_CAP_NUDGE, + parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, + strip_tool_markup as _shared_strip_tool_markup, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -48,12 +60,6 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs -from core.inference.tool_call_parser import ( - RAG_MAX_SEARCHES_PER_TURN, - RAG_SEARCH_CAP_NUDGE, - TOOL_XML_SIGNALS, - parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, -) from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, @@ -220,7 +226,7 @@ _INTENT_SIGNAL = re.compile( r"\b(?:now i|next i)\b" r")" ) -_MAX_REPROMPTS = 1 +_MAX_REPROMPTS = 3 # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -7881,12 +7887,17 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── @staticmethod - def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]: - """Thin wrapper around the shared parser in tool_call_parser - so safetensors and llama_cpp pick up the same fixes.""" + def _parse_tool_calls_from_text( + content: str, + *, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, + ) -> list[dict]: + """Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form.""" return _shared_parse_tool_calls_from_text( content, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) @staticmethod @@ -8406,11 +8417,17 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text - return strip_tool_call_markup(text, final = final) + return _shared_strip_tool_markup(text, final = final) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text + # Shared patterns so a textual Mistral/Llama call entering DRAINING is stripped, not + # leaked. Mistral first; no final trim so incremental length comparisons hold. + text = _strip_mistral_closed_calls(text) + # Parser-accurate function-XML scan before the regex arms so a literal ```` + # in a value doesn't make the tail eat trailing prose after the real ````. + text = _strip_function_xml_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -8456,6 +8473,13 @@ class LlamaCppBackend: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -8469,6 +8493,8 @@ class LlamaCppBackend: ) _MAX_BUFFER_CHARS = 32 + # Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors). + _MAX_BARE_JSON_BUFFER = 16384 _append_budget_exhausted_nudge = True # RAG: cap knowledge-base searches per assistant turn. The controller is # tool-agnostic, so this gate stays in the loop. @@ -8481,6 +8507,9 @@ class LlamaCppBackend: # "Hello!" won't match. Pattern compiled at module level # (_INTENT_SIGNAL). _reprompt_count = 0 + # Gates ``max_tool_iterations`` on real tool turns so reserved re-prompt slots don't + # extend the budget. Mirrors the safetensors guard. + _tool_iters_done = 0 _forced_tool_call_pending = False # Reserve extra iterations for re-prompts so they don't consume the @@ -8489,12 +8518,21 @@ class LlamaCppBackend: for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False active_tools = tool_controller.active_tools() if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS + # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + _enabled_tool_names = { + (tool.get("function") or {}).get("name") + for tool in active_tools + if (tool.get("function") or {}).get("name") + } + # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows. + _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -8777,7 +8815,36 @@ class LlamaCppBackend: is_prefix = True break - if is_match: + # Bare Llama-3.2 {"name":..} has no XML signal: hold an + # incomplete object, drain a complete one (mirrors safetensors). + _hold_buffer = False + # Whole buffer is the call (no visible prefix) -- drain silently. + _drain_silently = False + if not is_match and not is_prefix: + _bare = strip_llama3_leading_sentinels(stripped_buf) + if _bare.startswith("{"): + if _balanced_brace_end(_bare, 0) is None: + if len(stripped_buf) < _MAX_BARE_JSON_BUFFER: + _hold_buffer = True + elif _looks_like_enabled_bare_json( + _bare, _enabled_tool_names + ): + # Oversized still-open ENABLED-tool call: stop + # holding (memory bound) but DRAIN, not leak; + # a giant ordinary JSON answer still streams. + _drain_silently = True + elif self._parse_tool_calls_from_text( + content_buffer, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + _drain_silently = True + + if _drain_silently: + # No visible prefix -- the buffered text IS + # the call; drain without yielding it. + detect_state = _S_DRAINING + elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. @@ -8794,7 +8861,9 @@ class LlamaCppBackend: "text": cleaned, } detect_state = _S_DRAINING - elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + elif _hold_buffer or ( + is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS + ): pass # keep buffering else: # Not a tool -- flush buffer @@ -8821,8 +8890,16 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() + # A held bare-JSON fragment has no XML signal; route it to DRAINING. + _bare_eos = strip_llama3_leading_sentinels(stripped_buf) + # Gate on enabled names so a JSON answer isn't routed to DRAINING and dropped. + _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ) if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): detect_state = _S_DRAINING + elif _is_bare_tc: + detect_state = _S_DRAINING elif content_accum or reasoning_accum: detect_state = _S_STREAMING if content_buffer: @@ -8848,20 +8925,24 @@ class LlamaCppBackend: "text": cumulative_display, } else: + # No tool signal and no enabled bare-JSON call: a leading ``{`` is an ordinary + # JSON answer and must be shown; any other partial-markup prefix is dropped. + _held = strip_llama3_leading_sentinels(content_buffer.lstrip()) + if _held.startswith("{") and not _suppress_visible_output: + yield {"type": "content", "text": _held} return # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content. The - # route layer resets prev_text on tool_start, so post-tool - # synthesis streams correctly even if content was emitted - # before the tool XML. - _safety_tc = None - if any(s in content_accum for s in _tool_xml_signals): - _safety_tc = self._parse_tool_calls_from_text( - content_accum, - allow_incomplete = auto_heal_tool_calls, - ) + # Safety net: re-parse the full content for tool calls. The route layer resets + # prev_text on tool_start, so post-tool synthesis streams correctly even if + # content was emitted before the tool XML. Unconditional (not gated on + # _tool_xml_signals): bare-JSON and Gemma wrapper-less calls carry no signal. + _safety_tc = self._parse_tool_calls_from_text( + content_accum, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ) if not _safety_tc: # ── Re-prompt on plan-without-action ── # If the model described its intent (forward-looking @@ -8978,10 +9059,13 @@ class LlamaCppBackend: for i in sorted(tool_calls_acc) if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) ] or None - if not tool_calls and any(s in content_accum for s in _tool_xml_signals): + if not tool_calls: + # Unconditional re-parse: DRAINING means the buffer looked like a call, and + # bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if tool_calls and not has_structured_tc: content_text = _strip_tool_markup( @@ -8989,6 +9073,11 @@ class LlamaCppBackend: final = True, force = True, ) + # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call + # so the executed call isn't replayed as text or next-turn history. + content_text = strip_leading_bare_json_call( + content_text, _enabled_tool_names + ) if tool_calls: logger.info( f"Parsed {len(tool_calls)} tool call(s) from " @@ -9002,6 +9091,13 @@ class LlamaCppBackend: if content_accum: # Strip leaked tool-call XML before yielding. content_accum = _strip_tool_markup(content_accum, final = True) + # A truncated bare-JSON call has no XML to strip and didn't parse. With + # Auto-Heal on drop a leading ENABLED-tool fragment (plain JSON untouched); + # off keeps it visible per the strict contract. + if content_accum and active_tools and auto_heal_tool_calls: + content_accum = strip_leading_bare_json_call( + content_accum, _enabled_tool_names + ) if content_accum: yield {"type": "content", "text": content_accum} _meta = _build_metadata_event( @@ -9144,6 +9240,8 @@ class LlamaCppBackend: _kb_search_count += 1 completion = tool_controller.record_result(decision, result) resolved_provisional_tool_call_ids.add(decision.tool_call_id) + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -9167,6 +9265,12 @@ class LlamaCppBackend: if tool_controller.force_final_answer or not tool_controller.active_tools(): _append_budget_exhausted_nudge = False break + # Count only real tool turns against the cap so reserved re-prompt slots can't + # become extra tool rounds; a no-op turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _tool_iters_done += 1 + if _tool_iters_done >= max_tool_iterations: + break continue except httpx.ConnectError: diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index c73134b4a2..35855cc34d 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -29,10 +29,23 @@ import os from collections.abc import Mapping from typing import Any, Optional -from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal from core.inference.tool_loop_controller import coerce_tool_arguments from core.tool_healing import parse_tool_calls_from_text +# Only the formats this healer can promote. The parser's broader list adds Llama +# <|python_tag|> / Mistral [TOOL_CALLS], but buffering those here would flush a +# streamed call as prose, so keep a healer-aligned list. +_HEAL_SIGNALS = ( + "", + "<|tool_call>", + " bool: + return any(s in text for s in _HEAL_SIGNALS) + + # Read once at import (same convention as the other UNSLOTH_* switches). _HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" # Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process @@ -44,7 +57,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool: return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) -_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS) +_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS) # A suspected-but-unclosed tool block larger than this is declared a false # alarm and flushed, bounding memory on a model rambling XML-lookalike text. _MAX_HOLD_CHARS = 64 * 1024 @@ -198,7 +211,7 @@ def heal_openai_message_events( if not isinstance(msg, dict) or msg.get("tool_calls"): return None content = msg.get("content") - if not isinstance(content, str) or not has_tool_signal(content): + if not isinstance(content, str) or not _has_heal_signal(content): return None parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None @@ -248,7 +261,7 @@ def heal_openai_message( def _earliest_signal(buffer: str) -> int: best = -1 - for signal in TOOL_XML_SIGNALS: + for signal in _HEAL_SIGNALS: index = buffer.find(signal) if index >= 0 and (best < 0 or index < best): best = index @@ -275,7 +288,7 @@ def _partial_signal_suffix(buffer: str) -> int: """Length of the longest buffer suffix that is a proper prefix of a signal.""" for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): tail = buffer[-length:] - if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS): + if any(signal.startswith(tail) for signal in _HEAL_SIGNALS): return length return 0 @@ -508,7 +521,7 @@ def nudge_should_retry( if not message or message.get("tool_calls"): return False text = message.get("content") - if not isinstance(text, str) or not has_tool_signal(text): + if not isinstance(text, str) or not _has_heal_signal(text): return False return not _heal_would_promote(text, allowed_tools, tools) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 0c96378d6c..b67c6cf7e7 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -22,11 +22,17 @@ from loggers import get_logger from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_mistral_closed_calls, + _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, parse_tool_calls_from_text, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, strip_tool_markup, ) from core.inference.tool_loop_controller import ( @@ -50,6 +56,34 @@ logger = get_logger(__name__) # Buffer cap while disambiguating a possible tool-call prefix. _MAX_BUFFER_CHARS = 32 +# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. +_MAX_BARE_JSON_BUFFER = 16384 + +# Forward-looking intent ("I'll", "First,", "Step 1:") = planning; nudge a call. Negative +# lookahead drops negated forms ("I will not"). Mirrors GGUF. +_INTENT_SIGNAL = re.compile( + r"(?i)(" + r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" + r"|\b(?:first\b|step \d+:?|here['’]?s (?:my |the |a )?(?:plan|approach))" + r"|\b(?:now i|next i)\b" + r")" +) +_MAX_REPROMPTS = 3 +_REPROMPT_MAX_CHARS = 2000 +# Templated so the nudge names the caller's enabled tools. Mirrors GGUF tool_hint. +_REPROMPT_INSTRUCTION_TEMPLATE = ( + "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." +) + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + def strip_tool_markup_streaming( text: str, @@ -60,6 +94,12 @@ def strip_tool_markup_streaming( """Strip open-ended tool XML from display text without trimming whitespace.""" if not (auto_heal_tool_calls or tool_protocol_active): return text + # Mirror the final strip (no final trim): drop a leading Magistral ``[THINK]...[/THINK]`` + # block, then Mistral calls, then a parser-accurate function-XML scan before the regex + # arms. An unclosed ``[THINK]`` holds until ``[/THINK]`` so text stays monotonic. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + text = _strip_function_xml_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -81,6 +121,14 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + + _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') @@ -198,6 +246,10 @@ def run_safetensors_tool_loop( kb_search_count = 0 final_attempt_done = False next_call_id = 0 + reprompt_count = 0 + # Only turns that executed a tool count against ``max_tool_iterations``; a no-op or + # re-prompt turn must not consume budget (GGUF parity). + _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" @@ -215,9 +267,13 @@ def run_safetensors_tool_loop( _state_streaming = 1 _state_draining = 2 - for iteration in range(max_tool_iterations + 1): + # Reserve re-prompt slots so they don't eat the caller's tool budget. + _extra_iters = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False if final_attempt_done: active_tools: list[dict] = [] @@ -229,6 +285,8 @@ def run_safetensors_tool_loop( tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () + # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) detect_state = _state_buffering content_buffer = "" @@ -367,6 +425,34 @@ def run_safetensors_tool_loop( is_prefix = True break + # Bare Llama-3.2 ``{"name":..,"parameters":..}`` carries no XML signal. Hold a leading + # ``{`` (after any sentinel) until it closes: drain if it parses as a call, else stream. + bare_probe = strip_llama3_leading_sentinels(stripped) + if ( + not is_match + and not is_prefix + and tool_protocol_active + and bare_probe.startswith("{") + ): + if _balanced_brace_end(bare_probe, 0) is None: + if len(stripped) < _MAX_BARE_JSON_BUFFER: + continue # object still open -- keep buffering + elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names): + # Oversized still-open ENABLED-tool call: stop holding (memory bound) but + # DRAIN, not leak; a giant ordinary JSON answer still streams. + detect_state = _state_draining + continue + elif parse_tool_calls_from_text( + content_buffer, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + # Closed object that parses as a bare-JSON call -- drain silently. + detect_state = _state_draining + continue + # Closed non-call object (or oversized non-call) -- stream as text. + if is_match: # Tool signal -- flush any visible prefix before DRAINING # so the route sends it before tool_start. @@ -419,44 +505,74 @@ def run_safetensors_tool_loop( if detect_state == _state_buffering: # Buffer never resolved -- tool XML or plain content? stripped = content_buffer.lstrip() + _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active and any(sig in stripped for sig in tool_xml_signals) ): detect_state = _state_draining + elif tool_protocol_active and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ): + # Held ENABLED-tool bare-JSON fragment has no XML signal; DRAIN it (a JSON answer + # falls through to the else and streams, GGUF parity). + detect_state = _state_draining else: + # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser + # still fire on short emissions like "Let me search." that never exit BUFFERING. if content_buffer: cumulative_display += content_buffer - yield { - "type": "content", - "text": _strip_tool_markup_final( - cumulative_display, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } - yield {"type": "status", "text": ""} - return + cleaned = strip_tool_markup(cumulative_display, final = True) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + detect_state = _state_streaming if detect_state == _state_streaming: - # No tool detected mid-stream -- check for late tool XML. - safety_tc = None - saw_tool_signal = tool_protocol_active and any( - sig in content_accum for sig in tool_xml_signals + # Run the parser even with no XML signal (bare-JSON carries none); it's strict so + # plain answers stay untouched. Mirrors GGUF. + safety_tc = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) - if saw_tool_signal: - safety_tc = parse_tool_calls_from_text( - content_accum, - id_offset = next_call_id, - allow_incomplete = auto_heal_tool_calls, - ) if not safety_tc: - # Final answer: if a literal tool marker in prose was stripped - # during streaming but did not parse as a real call, restore the - # raw cumulative text for core callers. Route-level cleanup can - # still apply the Auto-Heal display policy. - if saw_tool_signal and content_accum: + # Re-prompt only when the model planned without acting (intent signal); + # "4" / "Hello!" never trigger. Mirrors GGUF. + _stripped = content_accum.strip() + if ( + tools + and auto_heal_tool_calls + and reprompt_count < _MAX_REPROMPTS + and 0 < len(_stripped) < _REPROMPT_MAX_CHARS + and _INTENT_SIGNAL.search(_stripped) + and not final_attempt_done + ): + reprompt_count += 1 + logger.info( + "Safetensors re-prompt %d/%d: model planned without " + "calling tools (%d chars)", + reprompt_count, + _MAX_REPROMPTS, + len(_stripped), + ) + tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" + conversation.append({"role": "assistant", "content": _stripped}) + conversation.append( + { + "role": "user", + "content": _REPROMPT_INSTRUCTION_TEMPLATE.format(tool_hint = tool_hint), + } + ) + yield {"type": "status", "text": ""} + continue + + # Final answer. If a literal tool marker in prose was buffered but never + # parsed as a call, restore the raw text so the prose surfaces; route + # cleanup still applies the Auto-Heal policy. + if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} yield {"type": "status", "text": ""} return @@ -476,20 +592,24 @@ def run_safetensors_tool_loop( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup # strips unparseable tool XML; disabled Auto-Heal preserves # the raw text so literal/malformed markup stays visible. if content_accum: - yield { - "type": "content", - "text": _strip_tool_markup_final( - content_accum, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } + _drain_text = _strip_tool_markup_final( + content_accum, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = False, + ) + # Drained bare-JSON call that didn't parse: with Auto-Heal on drop the fragment + # (plain JSON untouched); off keeps it visible per the strict contract. + if tool_protocol_active and auto_heal_tool_calls: + _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) + if _drain_text: + yield {"type": "content", "text": _drain_text} if provisional_render_html_started and not provisional_resolved: provisional_resolved = True yield { @@ -509,6 +629,9 @@ def run_safetensors_tool_loop( if tool_calls: next_call_id += len(tool_calls) + # Strip a leading bare-JSON call so it isn't replayed as text or next-turn history + # (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. + content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) if final_attempt_done: # Final-answer turn re-called a tool -- stop the loop. @@ -634,6 +757,8 @@ def run_safetensors_tool_loop( completion = tool_controller.record_result(decision, result) if provisional_match: provisional_resolved = True + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -646,7 +771,10 @@ def run_safetensors_tool_loop( if not unrestricted_tools and not tool_controller.active_tools(): final_attempt_done = True continue - if iteration + 1 >= max_tool_iterations and not final_attempt_done: + # Count only real tool turns against the cap so a no-op turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _executed_tool_iters += 1 + if _executed_tool_iters >= max_tool_iterations and not final_attempt_done: # Budget exhausted; nudge a final plain answer. final_attempt_done = True conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE}) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index ca3d1e4cbc..c31f4b272e 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,39 +2,74 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Backend-neutral tool-call XML parser shared by GGUF and safetensors. -Tolerates missing closing tags in either ``{json}`` -or ``v...`` shape. +Backend-neutral tool-call parser shared by GGUF, safetensors, and MLX, so the +safetensors + MLX agentic loop sees the same call shape llama-server gives GGUF: + + - ``{json}`` (Qwen / Hermes) + - ``v`` (Qwen3.5 xml) + - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) + - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) + - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) + - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + +Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ +# Keeps PEP 604 `X | None` lazy for python 3.9 (imported standalone by external servers). +from __future__ import annotations + +import json +import re +from typing import Any, Optional + +# Shared parser handles Qwen/Hermes, Qwen3.5 XML, Gemma 4; this module adds Llama-3, Mistral, bare JSON. from core import tool_healing as _tool_healing -_TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS +# Flip the streaming buffer STREAMING->DRAINING so partial markup never leaks. +TOOL_XML_SIGNALS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", +) -def parse_tool_calls_from_text( - content: str, - *, - id_offset: int = 0, - allow_incomplete: bool = True, -) -> list[dict]: - return _tool_healing.parse_tool_calls_from_text( - content, - id_offset = id_offset, - allow_incomplete = allow_incomplete, - ) +# Closed pairs only (mid-stream); _TOOL_ALL_PATS eats unclosed tails at end-of-turn. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + # Match to the real ```` (lookahead, not greedy ``.*``) so a literal + # ```` in a value doesn't truncate and each call stays separate. + re.compile( + r'' + r'(?:(?!).)*' + r"", + re.DOTALL, + ), + re.compile(r"<\|tool_call>.*?", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r'.*$', re.DOTALL), + # Bare-word markers drop a trailing truncated call only when the next chars look like + # a call start, so prose mentioning the marker is kept; a marker at end-of-text drops. + re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), + re.compile( + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\()|\s*$).*$", + re.DOTALL, + ), +] -def strip_tool_markup(text: str, *, final: bool = False) -> str: - return _tool_healing.strip_tool_call_markup(text, final = final) - - -# Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "<|tool_call>", "{json}``. +_TC_JSON_START_RE = re.compile(r"\s*\{") +# Qwen3.5 ```` plus attribute form ```` (MiniCPM-5, +# MiniMax-M2); name in group(1) or group(2). +_TC_FUNC_START_RE = re.compile(r'\s*') +# Body ends at ```` or ```` so trailing prose stays out of args. +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +# Horizontal whitespace only so the wrapping newline + indent survive (``_trim_param_value`` +# trims one newline), preserving code indent. +_TC_PARAM_START_RE = re.compile( + r'<(?:parameter|param)(?:=([\w\.\-]+)|\s+name="([\w\.\-]+)")>[^\S\n]*' +) +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + +# Llama-3 ``<|python_tag|>NAME.call(...)``. +_LLAMA3_PYTHON_TAG = "<|python_tag|>" +_LLAMA3_PY_CALL_RE = re.compile( + r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", +) +# Anchored at the char after ``<|python_tag|>`` plus the ``; NAME.call(`` chain sep, so +# a ``.call(`` inside JSON args is ignored. +_LLAMA3_PY_CALL_HEAD_RE = re.compile(r"\s*([\w\.\-]+)\s*\.\s*call\s*\(") +_LLAMA3_CALL_CHAIN_RE = re.compile(r"\s*;\s*([\w\.\-]+)\s*\.\s*call\s*\(") +# ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay linear on a +# truncated body (ReDoS). +_LLAMA3_KEY_RE = re.compile(r"\w+") +_LLAMA3_WS_RE = re.compile(r"\s*") +# ints, decimals, sci notation; trailing ``(?![\w.])`` stops ``1.2.3`` truncating to ``1.2``. +_LLAMA3_NUM_RE = re.compile(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?(?![\w.])") +_LLAMA3_LIT_RE = re.compile(r"true|false|null") + +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains ``name{json}`` (Magistral) or +# ``name[ARGS]{json}`` (Ministral / Large 3). +_MISTRAL_TRIGGER = "[TOOL_CALLS]" +_MISTRAL_ARGS_MARKER = "[ARGS]" +# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / Magistral). +_MISTRAL_CALL_ID_MARKER = "[CALL_ID]" +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside is not a real call. +_MISTRAL_THINK_OPEN = "[THINK]" +_MISTRAL_THINK_CLOSE = "[/THINK]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +_GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") +_GEMMA_STR_BEGIN = '<|"|>' +_GEMMA_STR_END = '<|"|>' +_GEMMA_TC_END = "" + + +def _balanced_bracket_end(text: str, start: int) -> int | None: + """Index of the ``]`` matching ``[`` at ``text[start]`` (ignores brackets in JSON strings).""" + if start >= len(text) or text[start] != "[": + return None + depth = 0 + in_string = False + esc = False + i = start + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _skip_mistral_call_id(text: str, pos: int) -> int: + """Skip an optional ``[CALL_ID]`` (Mistral Small 3.2); return the next token pos.""" + n = len(text) + i = pos + while i < n and text[i] in " \t\n\r": + i += 1 + if not text.startswith(_MISTRAL_CALL_ID_MARKER, i): + return pos + i += len(_MISTRAL_CALL_ID_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + # The id is a short opaque token; stop at whitespace or the next marker. + while i < n and text[i] not in " \t\n\r[{": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + return i + + +def _strip_mistral_reasoning(content: str) -> str: + """Drop a leading Magistral ``[THINK]`` block so rehearsed calls inside reasoning are not promoted; unclosed drops to EOF.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if not content.startswith(_MISTRAL_THINK_OPEN, i): + return content + close = content.find(_MISTRAL_THINK_CLOSE, i + len(_MISTRAL_THINK_OPEN)) + if close == -1: + return content[:i] + return content[:i] + content[close + len(_MISTRAL_THINK_CLOSE) :] + + +def _strip_mistral_closed_calls(text: str) -> str: + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks via balanced scanning (a non-greedy regex would truncate nested JSON); unclosed runs wait for ``final=True``.""" + n = len(text) + out = [] + cursor = 0 + while cursor < n: + idx = text.find(_MISTRAL_TRIGGER, cursor) + if idx == -1: + out.append(text[cursor:]) + break + out.append(text[cursor:idx]) + body_start = idx + len(_MISTRAL_TRIGGER) + i = body_start + while i < n and text[i] in " \t\n\r": + i += 1 + # Array shape: ``[TOOL_CALLS] [...]``. + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + if end is None: + # Truncated; let caller buffer / final-strip. + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Single-object shape ``[TOOL_CALLS] { json }``: the parser accepts it, so strip it too. + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Named shape: ``[TOOL_CALLS] name [ARGS]? { json }``. + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + out.append(text[idx:body_start]) + cursor = body_start + continue + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + out.append(text[idx:i]) + cursor = i + continue + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + # Consume the optional EOS marker so ``...{json}`` doesn't leave ```` as content. + if text.startswith("", cursor): + cursor += len("") + return "".join(out) + + +_FUNC_CLOSE_TAG_RE = re.compile(r"") + + +def _strip_function_xml_calls(text: str, *, final: bool) -> str: + """Strip ```` calls by mirroring the parser: an opener inside an open ```` is data and each call closes at its first ```` that is not parameter data; ``final`` drops a trailing unclosed call.""" + starts = [ + m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start()) + ] + if not starts: + return text + out: list[str] = [] + pos = 0 + for idx, m in enumerate(starts): + if m.start() < pos: + continue # opener already inside a previously consumed call span + out.append(text[pos : m.start()]) + next_start = starts[idx + 1].start() if idx + 1 < len(starts) else len(text) + close = None + for cm in _FUNC_CLOSE_TAG_RE.finditer(text, m.end(), next_start): + if not _inside_open_parameter(text, cm.start()): + close = cm # first close that is not parameter data = the real close + break + if close is not None: + pos = close.end() + elif final: + pos = len(text) # trailing unclosed call -- drop to EOF + else: + out.append(text[m.start() :]) # keep the unclosed call buffered mid-stream + pos = len(text) + break + out.append(text[pos:]) + return "".join(out) + + +def strip_tool_markup(text: str, *, final: bool = False) -> str: + """Strip tool-call markup; ``final=True`` also drops trailing unclosed runs and trims.""" + if final: + # End-of-turn only: drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket form, + # not the ```` reasoning channel) so raw reasoning doesn't leak into display/history. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + # Scan-strip the function-XML form first (parser-accurate: a literal ```` in + # a value is data, not a call); the regex arms below cover the other formats. + text = _strip_function_xml_calls(text, final = final) + pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + for pat in pats: + text = pat.sub("", text) + return text.strip() if final else text + + def has_tool_signal(text: str) -> bool: - """Return True if ``text`` contains any tool-call XML signal.""" return any(s in text for s in TOOL_XML_SIGNALS) + + +def _mistral_region_end(text: str, idx: int) -> int | None: + """Exclusive end of the balanced ``[TOOL_CALLS]`` call at ``idx``, or ``None`` when truncated (array, object, and named forms).""" + n = len(text) + i = idx + len(_MISTRAL_TRIGGER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + return None if end is None else end + 1 + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + return None + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + return None + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + + +def _xml_signal_inside_leading_mistral(content: str) -> bool: + """True when a parseable Mistral call is the first tool emission in document order: it owns the turn, so later XML (quoted in its arguments or in trailing prose) is not promoted over it. A signal BEFORE the trigger keeps normal order.""" + trig = content.find(_MISTRAL_TRIGGER) + if trig < 0: + return False + first_xml = _first_foreign_tool_signal(content) + if first_xml is not None and first_xml < trig: + return False + # Only plain prose precedes the trigger (preamble-tolerant); prose merely mentioning + # the marker has no parseable region and keeps the normal order. + return _mistral_region_end(content, trig) is not None + + +_ATTR_FUNC_OPEN_RE = re.compile(r' int | None: + """Offset of the first signal a non-envelope parser would fire on (XML forms plus the Llama-3 ``<|python_tag|>`` marker).""" + first = None + for sig in ("", "<|tool_call>", ""): + p = content.find(sig) + if p >= 0 and (first is None or p < first): + first = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first is None or attr.start() < first): + first = attr.start() + return first + + +def _xml_signal_inside_leading_bare_json(content: str) -> bool: + """True when the first foreign signal sits inside a LEADING bare-JSON call's balanced body: quoted argument data, so the bare-JSON parser takes the outer call first.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if i >= n or content[i] != "{": + return False + end = _balanced_brace_end(content, i) + if end is None: + return False + if _top_level_bare_json_name(content[i : end + 1]) is None: + # Not a call object, but a nameless object that parses as real JSON is an envelope + # too (markup in its strings is data); non-JSON braced prose keeps the old behaviour. + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a JSON envelope too, so fold it into first_xml. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + # Inside the balanced body the signal is quoted argument data, so the leading call owns + # the turn; a non-call object takes the decline path (dropped, only the tail parsed). + return first_xml is not None and i < first_xml + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Return OpenAI-format tool calls, first-match wins. ``allow_incomplete`` heals truncated calls (``False`` = strict closed-only); ``enabled_tool_names`` gates the markerless bare-JSON form.""" + # Drop Magistral reasoning before any dispatch so a rehearsed call inside + # [THINK]...[/THINK] is not promoted; keeps the parse path aligned with the display strip. + content = _strip_mistral_reasoning(content) + + # A leading bare-JSON value is decided FIRST so markup quoted in its arguments stays + # data. Must precede the Mistral guard, whose preamble tolerance would else claim a + # trigger quoted inside the leading object. + if _xml_signal_inside_leading_bare_json(content): + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + # Disabled/example name: the leading object is ordinary content. Drop it and parse + # only the tail -- a real call after it still parses, nothing inside it is promoted. + i = 0 + while i < len(content) and content[i] in " \t\n\r": + i += 1 + end = _balanced_brace_end(content, i) # guard guarantees a balanced object + return parse_tool_calls_from_text( + content[end + 1 :], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A [TOOL_CALLS] call that is the first tool emission owns the turn: XML quoted in its + # arguments or in trailing prose is not promoted, and a plain-prose preface keeps it. + if _xml_signal_inside_leading_mistral(content): + calls = _parse_mistral_tool_calls( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading MiniCPM/MiniMax ```` call owns the turn: tool_healing + # does not know the wrapper, so gate it here. A signal before the opener keeps normal order. + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None: + first_other = None + for sig in ( + "", + "<|tool_call>", + "", + _MISTRAL_TRIGGER, + ): + p = content.find(sig) + if p >= 0 and (first_other is None or p < first_other): + first_other = p + if first_other is None or attr.start() < first_other: + calls = _parse_function_xml( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading Llama-3 ``<|python_tag|>`` call owns the turn like the others: markup quoted + # in a ``.call(...)`` argument is not promoted. tool_healing does not know the tag, so + # gate it here. A foreign signal before the tag keeps normal order. + py_tag = content.find(_LLAMA3_PYTHON_TAG) + if py_tag >= 0: + first_other = None + for sig in ("", "<|tool_call>", "= 0 and (first_other is None or p < first_other): + first_other = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first_other is None or attr.start() < first_other): + first_other = attr.start() + if first_other is None or py_tag < first_other: + calls = _parse_llama3_python_tag( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # Qwen/Hermes, Qwen3.5 XML, and Gemma 4 use the shared tool_healing parser (the + # strict/Auto-Heal + nested-marker + ``<|"|>`` handling GGUF relies on). + calls = _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + ) + if calls: + return calls + + # Formats tool_healing does not cover: ```` (MiniCPM-5 / MiniMax-M2), + # Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected + # call is never re-healed here. + for parser in ( + _parse_function_xml, # attribute form + _parse_llama3_python_tag, # Llama-3 <|python_tag|> + _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] + ): + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # Llama-3.2 bare ``{"name":..., "parameters":...}``. Strict (starts with ``{`` + # and parses to the right shape) so plain prose stays untouched. + return _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + + +def _parse_tool_call_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + for m in _TC_JSON_START_RE.finditer(content): + brace_start = m.end() - 1 + end = _balanced_brace_end(content, brace_start) + if end is None: + continue + # Strict mode: a balanced body that never closed its ```` is truncated + # (trailing prose after the close is still tolerated). + if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): + continue + try: + obj = json.loads(content[brace_start : end + 1]) + except (json.JSONDecodeError, ValueError): + continue + name = obj.get("name", "") + # Accept both ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if not name: + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + return out + + +def _trim_param_value(val: str) -> str: + """Trim only the template's wrapping newline around an XML parameter value; ``str.strip()`` destroyed code/diff indentation.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _inside_open_parameter(text: str, pos: int) -> bool: + """True if ``pos`` is inside an unclosed ```` block, i.e. the opener at ``pos`` is literal argument data, not a nested call.""" + last_param_open = -1 + for m in _TC_PARAM_START_RE.finditer(text, 0, pos): + last_param_open = m.start() + if last_param_open < 0: + return False + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal ````); an unclosed one falls back to func close. + own_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + if own_closes: + return min(own_closes) > pos + func_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + return not func_closes or pos < min(func_closes) + + +def _parse_function_xml( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + # Skip ```` openers that are literals inside an open parameter value, + # else the nested marker becomes a second call and truncates the real argument. + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + # group(1) is ````, group(2) is ````. + func_name = fm.group(1) or fm.group(2) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + # The call ends at the FIRST / not inside an open parameter: + # a literal close in an argument is skipped as data, prose after the real close is not + # folded in (mirrors _strip_function_xml_calls). + close_match = None + for cm in _TC_END_TAG_RE.finditer(content, body_start, next_func): + if not _inside_open_parameter(content, cm.start()): + close_match = cm + break + has_close = close_match is not None + if has_close: + body_end = close_match.start() + else: + body_end = min(len(content), next_func) + # Strict mode: a call that never reached its close is truncated; do not heal it. + if not allow_incomplete and not has_close: + continue + body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) + + args: dict = {} + param_unclosed = False + # Same nested-literal guard: a ```` opener inside an open value is literal text. + param_starts = [ + pm + for pm in _TC_PARAM_START_RE.finditer(body) + if not _inside_open_parameter(body, pm.start()) + ] + if len(param_starts) == 1: + pm = param_starts[0] + raw_val = body[pm.end() :] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + else: + for pidx, pm in enumerate(param_starts): + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + raw_val = body[val_start:next_param] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + + # Strict mode: every parameter must close; a dangling one means the call was cut off. + # A closed call with no parameters is a valid zero-argument call, so keep it. + if not allow_incomplete and param_unclosed: + continue + + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": func_name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: + """One ``.call`` value at ``body[p:]``; returns ``(value, len)`` or ``(None, None)``.""" + if p >= n: + return None, None + if body[p] == '"': + # ``"((?:\\.|[^"\\])*)"`` by hand so an unterminated quote is O(n), not O(n^2). + j = p + 1 + while j < n: + c = body[j] + if c == "\\": + # ``\\.`` needs a following non-newline char; else the body can't match. + if j + 1 >= n or body[j + 1] == "\n": + return None, None + j += 2 + continue + if c == '"': + raw = body[p + 1 : j] + # json.loads keeps \n/\uXXXX escapes and literal UTF-8 (emoji/CJK) intact. + try: + return json.loads('"' + raw + '"'), j + 1 - p + except (json.JSONDecodeError, ValueError): + return raw, j + 1 - p + j += 1 + return None, None # unterminated + nm = _LLAMA3_NUM_RE.match(body, p) + if nm: + v = nm.group(0) + # Sci notation and decimals decode as float; a bare integer stays int. + return (float(v) if any(c in v for c in ".eE") else int(v)), nm.end() - p + lm = _LLAMA3_LIT_RE.match(body, p) + if lm: + return {"true": True, "false": False, "null": None}[lm.group(0)], lm.end() - p + return None, None + + +def _parse_llama3_kv_args(body: str) -> dict[str, Any]: + """Left-to-right ``k=v`` kwargs from a ``.call(...)`` body (linear scan; later keys win).""" + args: dict[str, Any] = {} + n = len(body) + i = 0 + while i < n: + km = _LLAMA3_KEY_RE.match(body, i) + if km is None: + i += 1 + continue + p = _LLAMA3_WS_RE.match(body, km.end()).end() + if p >= n or body[p] != "=": + i = km.end() + continue + p = _LLAMA3_WS_RE.match(body, p + 1).end() + val, length = _llama3_kv_value(body, p, n) + if length is None: + i = km.end() + continue + args[km.group(0)] = val + i = p + length + return args + + +def _parse_llama3_python_tag( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse Llama-3 ``<|python_tag|>`` emissions: ``NAME.call(...)``, bare JSON, ``; `` multi-call, ``parameters``/``arguments`` keys.""" + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` (optionally + # ``; ``-chained) so a ``.call(...)`` inside a JSON string argument isn't mistaken for one. + pos = content.find(_LLAMA3_PYTHON_TAG) + truncated = False + while pos >= 0 and not truncated: + head = _LLAMA3_PY_CALL_HEAD_RE.match(content, pos + len(_LLAMA3_PYTHON_TAG)) + if head is None: + # Tag is the custom JSON form (``{...}``) or noise -- leave it to step 2. + break + name = head.group(1) + open_idx = head.end() + i = open_idx + while True: + i = open_idx + depth = 1 + in_string = False + esc = False + while i < len(content) and depth > 0: + ch = content[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + i += 1 + # Truncated ``.call(...)`` (no closing paren): reject in strict mode. + if not allow_incomplete and depth > 0: + truncated = True + break + body = content[open_idx:i] + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(_parse_llama3_kv_args(body)), + }, + } + ) + # ``)`` then optional ``; NAME.call(`` chains the next built-in call. + chain = _LLAMA3_CALL_CHAIN_RE.match(content, i + 1) + if chain is None: + break + name = chain.group(1) + open_idx = chain.end() + # Past the consumed region: a second ``<|python_tag|>`` may carry more calls. + pos = content.find(_LLAMA3_PYTHON_TAG, i + 1) + + # 2. ``<|python_tag|>{"name":.., "parameters":..}``; raw_decode peels ``; ``-separated objects. + if not out: + decoder = json.JSONDecoder() + idx = content.find(_LLAMA3_PYTHON_TAG) + while idx >= 0: + search_from = idx + len(_LLAMA3_PYTHON_TAG) + cursor = search_from + while cursor < len(content): + brace = content.find("{", cursor) + if brace < 0: + break + # Stop at the next ``<|python_tag|>``. + next_tag = content.find(_LLAMA3_PYTHON_TAG, search_from, brace) + if next_tag >= 0: + break + try: + obj, end_offset = decoder.raw_decode(content[brace:]) + except (json.JSONDecodeError, ValueError): + cursor = brace + 1 + continue + if not isinstance(obj, dict): + cursor = brace + end_offset + continue + name = obj.get("name") or obj.get("function") or "" + args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor = brace + end_offset + idx = content.find(_LLAMA3_PYTHON_TAG, cursor) + return out + + +# Llama-3 special-token sentinels (chainable, any order) plus the header role label. +_LLAMA3_BARE_JSON_SENTINELS = ( + "<|begin_of_text|>", + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", +) +_LLAMA3_HEADER_ROLES = ("assistant", "user", "system", "tool", "ipython") + + +def strip_llama3_leading_sentinels(content: str) -> str: + """Strip leading Llama-3 sentinels leaked from a prior turn; shared by the parser and the streaming guards.""" + stripped = content.lstrip() + while True: + stripped = stripped.lstrip() + matched = False + for sentinel in _LLAMA3_BARE_JSON_SENTINELS: + if stripped.startswith(sentinel): + stripped = stripped[len(sentinel) :] + if sentinel == "<|start_header_id|>": + for role in _LLAMA3_HEADER_ROLES: + if stripped.startswith(role): + stripped = stripped[len(role) :] + break + matched = True + break + if not matched: + return stripped + + +def _parse_llama3_bare_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Llama-3.2 bare ``{"name":.., "parameters":..}`` (strict). ``enabled_tool_names`` keeps ordinary JSON answers from being misread; ``None`` is name-agnostic.""" + out: list[dict] = [] + stripped = strip_llama3_leading_sentinels(content) + if not stripped.startswith("{"): + return out + + decoder = json.JSONDecoder() + cursor = 0 + n = len(stripped) + while cursor < n: + # Skip whitespace and the Llama-3 ``;`` inter-call separator. + while cursor < n and stripped[cursor] in " \t\n\r;": + cursor += 1 + if cursor >= n or stripped[cursor] != "{": + break + try: + obj, end_offset = decoder.raw_decode(stripped[cursor:]) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(obj, dict): + break + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + break + # Markerless JSON is ambiguous: only a call when the name is an enabled tool. + if enabled_tool_names is not None and name not in enabled_tool_names: + break + # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or a + # JSON-string of one (OpenAI). + if "parameters" in obj: + args = obj.get("parameters") + if not isinstance(args, dict): + break + args_str = json.dumps(args) + elif "arguments" in obj: + args = obj.get("arguments") + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(parsed, dict): + break + args_str = args + else: + break + else: + break + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor += end_offset + return out + + +def _parse_mistral_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse Mistral ``[TOOL_CALLS]`` emissions: pre-v11 array/object and v11+ named forms.""" + out: list[dict] = [] + content = _strip_mistral_reasoning(content) + idx = content.find(_MISTRAL_TRIGGER) + if idx < 0: + return out + + # Disambiguate the first occurrence: array / single object (pre-v11) or bare-name (v11+). + j = idx + len(_MISTRAL_TRIGGER) + k = j + while k < len(content) and content[k] in " \t\n\r": + k += 1 + if k >= len(content): + return out + + if content[k] == "[": + return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) + + if content[k] == "{": + # Pre-v11 single ``{"name":...}``; fall through to v11+ if it carries no ``name``. + end = _balanced_brace_end(content, k) + if end is not None: + try: + obj = json.loads(content[k : end + 1]) + if isinstance(obj, dict) and obj.get("name"): + _consume_mistral_call(content[k : end + 1], out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or ``name[ARGS]{json}``. + pos = idx + while pos >= 0: + cur = pos + len(_MISTRAL_TRIGGER) + nm = _MISTRAL_V11_NAME_RE.match(content, cur) + if not nm: + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + name = nm.group(1) + after_name = nm.end() + after_name = _skip_mistral_call_id(content, after_name) + if content.startswith(_MISTRAL_ARGS_MARKER, after_name): + after_name += len(_MISTRAL_ARGS_MARKER) + while after_name < len(content) and content[after_name] in " \t\n\r": + after_name += 1 + if after_name >= len(content) or content[after_name] != "{": + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + end = _balanced_brace_end(content, after_name) + if end is None: + break + try: + args = json.loads(content[after_name : end + 1]) + except (json.JSONDecodeError, ValueError): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + if not isinstance(args, dict): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = content.find(_MISTRAL_TRIGGER, end + 1) + return out + + +def _parse_mistral_array( + content: str, + start: int, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Pre-v11 ``[TOOL_CALLS] [{...}, ...]`` array form.""" + out: list[dict] = [] + j = start + depth = 0 + in_string = False + esc = False + while j < len(content): + ch = content[j] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + # An unclosed array (no matching ]) is truncated; reject in strict mode. + if not allow_incomplete and depth != 0: + return out + body = content[start : j + 1] if depth == 0 else content[start:] + + try: + arr = json.loads(body) + if isinstance(arr, list): + for obj in arr: + if isinstance(obj, dict): + _consume_mistral_call(json.dumps(obj), out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + if not allow_incomplete: + return out + + # Healing path for unclosed arrays: walk top-level objects, advancing past each + # balanced ``{...}`` (re-scanning from every ``{`` would be quadratic ReDoS). + pos = 0 + blen = len(body) + while pos < blen: + brace = body.find("{", pos) + if brace < 0: + break + end = _balanced_brace_end(body, brace) + if end is None: + break # truncated mid-object: nothing after it can balance + _consume_mistral_call(body[brace : end + 1], out, id_offset) + pos = end + 1 + return out + + +def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> None: + try: + obj = json.loads(obj_text) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(obj, dict): + return + name = obj.get("name") or "" + # Mistral uses ``arguments``; accept the ``parameters`` alias too. + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": obj.get("id") or f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + + +def _parse_gemma_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``.""" + out: list[dict] = [] + for m in _GEMMA_TC_RE.finditer(content): + name = m.group(1) + body_start = m.end() - 1 + end_marker = content.find(_GEMMA_TC_END, body_start) + # No closing tag: truncated call, reject in strict mode. + if not allow_incomplete and end_marker < 0: + continue + scan_end = end_marker if end_marker >= 0 else len(content) + end = _gemma_balanced_brace_end(content, body_start, scan_end) + if end is None: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_mapping_body(body) + except Exception: + args = {} + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _balanced_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` matching ``{`` at ``brace_pos`` (ignores braces in JSON strings).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + in_string = False + esc = False + i = brace_pos + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +_BARE_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"') + + +def _top_level_bare_json_name(probe: str) -> Optional[str]: + """Top-level ``"name"`` (or ``"function"`` alias) of a bare-JSON object, else None; nested objects are skipped and truncated tails return None.""" + if not probe.startswith("{"): + return None + decoder = json.JSONDecoder() + function_value = None # the ``"function"`` alias, used only if no ``"name"`` key + i = 1 + n = len(probe) + while i < n: + while i < n and probe[i] in " \t\r\n,": + i += 1 + if i >= n or probe[i] == "}": + # End of object, no top-level ``"name"``: fall back to the ``"function"`` alias. + return function_value + if probe[i] != '"': + return None + try: + key, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(key, str): + return None + i += consumed + while i < n and probe[i] in " \t\r\n": + i += 1 + if i >= n or probe[i] != ":": + return None + i += 1 + while i < n and probe[i] in " \t\r\n": + i += 1 + if key == "name": + if i < n and probe[i] == '"': + try: + value, _consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + return value if isinstance(value, str) else None + return None + if key == "function" and function_value is None and i < n and probe[i] == '"': + # ``"function"`` is an alias; record it but keep scanning (``"name"`` wins). + try: + value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if isinstance(value, str): + function_value = value + i += consumed + continue + # Skip a non-name top-level value; a truncated one returns None (keep the text). + if i < n and probe[i] == "{": + end = _balanced_brace_end(probe, i) + if end is None: + return None + i = end + 1 + elif i < n and probe[i] == "[": + end = _balanced_bracket_end(probe, i) + if end is None: + return None + i = end + 1 + else: + try: + _value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + i += consumed + # No top-level ``"name"`` key: fall back to the ``"function"`` alias if seen. + return function_value + + +def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run) + that ``strip_tool_markup`` misses; non-call text is unchanged and + ``enabled_tool_names`` gates like the parser. Consuming the whole chain + matters because the loops keep this text as next-turn assistant history: a + leftover executed call would be replayed alongside the structured + ``tool_calls``.""" + remainder = text + stripped_any = False + while True: + probe = strip_llama3_leading_sentinels(remainder.lstrip()) + # Skip the Llama-3 ``;`` inter-call separator between chained calls. + if stripped_any: + probe = probe.lstrip(" \t\n\r;") + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return probe.lstrip() if stripped_any else text + if enabled_tool_names is not None: + # Only suppress when the leading object's TOP-LEVEL name is an enabled tool + # (a nested ``"name"`` is data); an unknown name is kept. + name = _top_level_bare_json_name(probe) + if name not in enabled_tool_names: + return probe.lstrip() if stripped_any else text + end = _balanced_brace_end(probe, 0) + if end is None: + return "" # truncated bare-JSON call -- nothing recoverable + # A closed object must have the CALL SHAPE the parser accepts; an ordinary JSON + # answer it rejects is content, so keep it visible. + try: + obj = json.loads(probe[: end + 1]) + except (json.JSONDecodeError, ValueError): + return probe.lstrip() if stripped_any else text + if not _bare_json_call_shaped(obj): + return probe.lstrip() if stripped_any else text + remainder = probe[end + 1 :] + stripped_any = True + + +def _bare_json_call_shaped(obj) -> bool: + """The shape gate ``_parse_llama3_bare_json`` applies to a decoded object.""" + if not isinstance(obj, dict): + return False + # The parser requires a TOP-LEVEL name; a nested one is data, not the call name. + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + return False + if "parameters" in obj: + return isinstance(obj.get("parameters"), dict) + args = obj.get("arguments") + if isinstance(args, dict): + return True + if isinstance(args, str): + try: + return isinstance(json.loads(args), dict) + except (json.JSONDecodeError, ValueError): + return False + return False + + +def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int | None: + """Like ``_balanced_brace_end`` but skips ``<|"|>`` strings and matches {}/[] symmetrically.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + i = brace_pos + while i < hard_stop: + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return None + i = close + len(_GEMMA_STR_END) + continue + ch = text[i] + if ch == "{" or ch == "[": + depth += 1 + elif ch == "}" or ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_parse_value(text: str, i: int): + """Parse one Gemma arg value at ``i``; returns ``(value, next_index)``.""" + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return text[i + len(_GEMMA_STR_BEGIN) :], len(text) + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END) + if text[i] == "{": + end = _gemma_balanced_brace_end(text, i, len(text)) + if end is None: + return {}, len(text) + return _gemma_parse_mapping_body(text[i + 1 : end]), end + 1 + if text[i] == "[": + j, depth = i, 0 + while j < len(text): + if text.startswith(_GEMMA_STR_BEGIN, j): + k = text.find(_GEMMA_STR_END, j + len(_GEMMA_STR_BEGIN)) + if k < 0: + j = len(text) + break + j = k + len(_GEMMA_STR_END) + continue + ch = text[j] + if ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + body = text[i + 1 : j] + items: list[Any] = [] + k = 0 + while k < len(body): + if body[k] in " \t\n\r,": + k += 1 + continue + v, k = _gemma_parse_value(body, k) + items.append(v) + return items, j + 1 + # Primitive: number / true/false/null / bare identifier. + end = i + while end < len(text) and text[end] not in ",}]" and not text.startswith(_GEMMA_STR_BEGIN, end): + end += 1 + if end == i: + # Stray delimiter, nothing consumed: advance past it so callers can't spin forever. + return "", i + 1 + raw = text[i:end].strip() + if raw == "true": + return True, end + if raw == "false": + return False, end + if raw == "null": + return None, end + try: + return int(raw), end + except ValueError: + pass + try: + return float(raw), end + except ValueError: + pass + return raw, end + + +def _gemma_parse_mapping_body(body: str) -> dict[str, Any]: + """Parse a Gemma argument mapping (content between `{` and `}`).""" + out: dict[str, Any] = {} + i = 0 + n = len(body) + while i < n: + while i < n and body[i] in " \t\n\r,": + i += 1 + if i >= n: + break + if body.startswith(_GEMMA_STR_BEGIN, i): + close = body.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = body[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and body[i] != ":": + i += 1 + key = body[kstart:i].strip() + while i < n and body[i] in " \t\n\r": + i += 1 + if i < n and body[i] == ":": + i += 1 + while i < n and body[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + v, i = _gemma_parse_value(body, i) + out[key] = v + return out diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e8367ad08c..ff8faf2308 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -27,12 +27,15 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") +# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around +# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``. +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") +# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline). +_TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" @@ -43,7 +46,8 @@ _FUNC_CLOSE_TAG = "" # must be identifier-shaped (start with a letter or underscore); a comma # followed by digits-then-colon is value text such as a timestamp or ratio # (`meet at 10:00, 11:00 tomorrow`), not a new key. -_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") +# Dots match the key-quoting scanner: a dotted key after a bare value must end the value at the comma. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") def _balanced_brace_end( @@ -223,7 +227,9 @@ def _quote_gemma_object_keys(src: str) -> str: 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 "_-"): + # Dots match the parser's key/name charset: Gemma emits dotted argument keys + # (user.name:...) for namespaced schemas. + while i < len(src) and (src[i].isalnum() or src[i] in "_-."): i += 1 key_name = src[key_name_start:i] colon_pos = i @@ -267,7 +273,8 @@ def _quote_gemma_object_keys(src: str) -> str: json.loads(raw.strip()) parts.append(raw) except (json.JSONDecodeError, ValueError): - parts.append(json.dumps(raw.strip()) if raw.strip() else raw) + # Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}. + parts.append(json.dumps(raw.strip())) else: parts.append(src[key_start:i]) return "".join(parts) @@ -291,9 +298,35 @@ def _inside_open_parameter(content: str, pos: int) -> bool: last_param_start = match.start() if last_param_start < 0: return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal function closes); an unclosed one falls back to func close. + own_close = content.find(_PARAM_CLOSE_TAG, last_param_start) + if own_close >= 0: + return own_close > pos + func_close = content.find(_FUNC_CLOSE_TAG, last_param_start) + return func_close < 0 or pos < func_close + + +def _func_close_index(content: str, body_start: int, body: str) -> int: + """Index in ``body`` of the first ```` that is not argument + data (not inside an open parameter value); -1 when every close is data. + Taking the LAST close swallowed prose between the real close and a + literal ```` mentioned later in the answer.""" + idx = body.find(_FUNC_CLOSE_TAG) + while idx >= 0: + if not _inside_open_parameter(content, body_start + idx): + return idx + idx = body.find(_FUNC_CLOSE_TAG, idx + 1) + return -1 + + +def _trim_param_value(val: str) -> str: + """Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val def parse_tool_calls_from_text( @@ -349,7 +382,10 @@ def parse_tool_calls_from_text( if kind == "json": obj = json.loads(content[m.end() - 1 : end + 1]) name = obj.get("name", "") - arguments = obj.get("arguments", {}) + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). + arguments = obj.get("arguments") + if arguments is None: + arguments = obj.get("parameters", {}) if isinstance(arguments, dict): arguments = json.dumps(arguments) else: @@ -382,7 +418,7 @@ def parse_tool_calls_from_text( body_end = len(content) body_end = min(body_end, next_func) body = content[body_start:body_end] - close_idx = body.rfind(_FUNC_CLOSE_TAG) + close_idx = _func_close_index(content, body_start, body) if close_idx >= 0: span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) body = body[:close_idx] @@ -404,7 +440,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() + arguments[pm.group(1)] = _trim_param_value(val) else: valid_params = True for pidx, pm in enumerate(param_starts): @@ -422,7 +458,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() + arguments[param_name] = _trim_param_value(val) if not valid_params: continue @@ -444,6 +480,86 @@ def parse_tool_calls_from_text( } ) call_spans.append((start, span_end)) + + if not tool_calls: + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + # Span for with_spans callers: through the close if present, else body end. + span_end = body_end + if not allow_incomplete: + close_idx = _func_close_index(content, body_start, body) + if close_idx < 0: + continue + body = body[:close_idx] + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + else: + # Terminate at the real close so trailing prose doesn't leak in; no close -> whole body. + close_idx = _func_close_index(content, body_start, body) + if close_idx >= 0: + body = body[:close_idx] + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = _trim_param_value(val) + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = _trim_param_value(val) + if not valid_params: + continue + + tc = { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": func_name, + "arguments": json.dumps(arguments), + }, + } + tool_calls.append(tc) + call_spans.append((fm.start(), span_end)) + if with_spans: return tool_calls, call_spans return tool_calls diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17be222d93..4393c1b304 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -603,6 +603,17 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: ) +def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(reasoning_content = text), + finish_reason = None, + ) + + def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: """Terminal stop chunk (empty delta) carrying the finish reason.""" return _chat_chunk_sse( @@ -1136,6 +1147,7 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls from core.inference.passthrough_healing import ( StreamToolCallHealer, heal_gate, @@ -1294,6 +1306,11 @@ async def artifact_preview_frame(allow_network: bool = False): ) +# Whitespace/escape-tolerant bare-JSON tool-template detector: matches pretty-printed and +# JSON-escaped ``{"name":`` plus the ``"function"`` alias. +_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and @@ -1304,17 +1321,21 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: model_identifier = model_id, log_source = "safetensors", ) - # Our safetensors loop only parses {json}, - # ..., and Gemma native <|tool_call>.... - # Llama uses <|python_tag|>, Mistral uses [TOOL_CALLS]; advertising tools for - # those enables a pill the parser can't honour. GGUF is unaffected -- - # llama-server normalises every format into structured deltas. + # Markers the parser recognises; drop the pill if a template advertises tools but uses none. + # The bare-JSON ``{"name":`` form is matched whitespace-tolerantly below. + _PARSER_MARKERS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ) if ( flags.get("supports_tools") and chat_template - and "" not in chat_template - and "" not in chat_template + and not any(m in chat_template for m in _PARSER_MARKERS) + and not _BARE_JSON_NAME_MARKER_RE.search(chat_template) ): logger.info( "safetensors: template advertises tools but uses an " @@ -1335,6 +1356,31 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: return flags +def _sf_reasoning_prefill_mode( + features: dict, + enable_thinking: Optional[bool], + template: Optional[str] = None, + reasoning_effort: Optional[str] = None, +) -> bool: + """Whether this request begins inside an unclosed ```` (Qwen3/GLM prefill it). Gated on the standard markers; bespoke channels, gpt-oss, and thinking-disabled requests are excluded. ``enable_thinking=None`` defaults ON.""" + if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): + return False + tpl = template or "" + if "" not in tpl and "" not in tpl: + return False + if features.get("reasoning_always_on"): + return True + if not features.get("supports_reasoning"): + return False + if enable_thinking is False: + return False + # reasoning_effort="none" disables thinking on enable_thinking_effort (GLM-5.2) models like + # enable_thinking=False; without this the answer is swallowed into empty reasoning_content. + if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": + return False + return True + + def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. @@ -1605,30 +1651,41 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py -# split across the visible/DRAIN boundary. Four leak shapes: -# 1. well-formed `...` / `...` -# 2. orphan opening to EOF (close was DRAINED) -# 3. bare orphan close (open was DRAINED) -# 4. tail-only `` (outer close truncated by EOS); anchored to -# `\Z` so mid-text `` in user code samples survives. +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# ``llama_cpp.py``'s speculative buffer splits across the visible/DRAIN boundary. Mistral +# ``[TOOL_CALLS]`` uses the parser's balanced-brace helper (``\{.*?\}`` would truncate nested JSON). _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-]+)>.*?(?:|\Z)" + # The ``<|python_tag|>`` arm runs to the next REAL Llama sentinel or EOF, so a literal + # ``<|...|>`` token in an argument (e.g. ``<|cite|>``) doesn't truncate the strip. + # ```` plus the ```` attribute form; name class mirrors the parser. + # A CLOSED ``...`` extends to the last ```` before the next + # opener (so a literal ```` in a value can't truncate); this arm runs first. + r'(?:(?!).)*' + r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:|\Z)' r"|<\|tool_call>.*?(?:|\Z)" r"|" r"|" - r"|\s*\Z", + r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + # ```` is the attribute-form alias of ````; strip a tail-only orphan. + r"|\s*\Z", _re.DOTALL, ) +def _strip_tool_xml(text: str) -> str: + """Mistral balanced-brace helper + guarded function-XML scan + ``_TOOL_XML_RE`` (skips openers inside an open ````).""" + return _TOOL_XML_RE.sub( + "", _strip_function_xml_calls(_strip_mistral_closed_calls(text), final = True) + ) + + def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: - """Apply route-level XML leak cleanup only when Auto-Heal is enabled.""" + """Route-level tool-call leak cleanup (Auto-Heal only) via ``_strip_tool_xml``.""" if not auto_heal_tool_calls: return text - return _TOOL_XML_RE.sub("", text) + return _strip_tool_xml(text) logger = get_logger(__name__) @@ -6511,6 +6568,22 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) + # Split prefilled-```` output into reasoning_content deltas (GGUF parity) so the UI + # renders the thinking block for safetensors and MLX. + _sf_parse_think = bool( + _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") + ) + # Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded. + _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( + _sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort + ) + + def _new_sf_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -6652,6 +6725,19 @@ async def openai_chat_completions( gen = sf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() + + def _sf_flush_reasoning(): + # Drain the extractor at a turn boundary / stream end; only visible text reaches the monitor. + fr, fv = reasoning_extractor.finish() + out = [] + if fr: + out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr)) + if fv: + api_monitor.append_reply(monitor_id, fv) + out.append(_chat_content_chunk(completion_id, created, model_name, fv)) + return out + while True: if cancel_event.is_set(): backend.reset_generation_state() @@ -6668,7 +6754,11 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: + # Turn boundary: flush reasoning, then start a fresh extractor. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() status_data = json.dumps( { "type": "tool_status", @@ -6680,7 +6770,11 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + # Flush reasoning before tool_start so the thinking block closes ahead of the tool card. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -6694,9 +6788,18 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split reasoning vs visible; only visible reaches the monitor. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + for _c in _sf_flush_reasoning(): + yield _c yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk from the last turn, same shape as the # GGUF tool loop's metadata. Request-scoped holder, so @@ -6774,18 +6877,27 @@ async def openai_chat_completions( return full_text content_text = await asyncio.to_thread(_drain_to_text) - api_monitor.set_reply(monitor_id, content_text) + # Split prefilled reasoning from the visible answer; monitor gets visible text only. + _reasoning_text, _visible_text = _extract_responses_reasoning( + content_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + api_monitor.set_reply(monitor_id, _visible_text) _stats = _sf_stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") + _sf_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _sf_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = content_text), + message = CompletionMessage(**_sf_msg_kwargs), finish_reason = "stop", ) ], @@ -6864,6 +6976,8 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" + # Split prefilled into reasoning_content deltas. Single turn (no per-turn reset); also MLX. + reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive # concurrently but the orchestrator serializes them via @@ -6892,9 +7006,21 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. @@ -6956,18 +7082,27 @@ async def openai_chat_completions( for token in generate(): full_text = token + # Split prefilled reasoning from the visible answer; also covers MLX. + _reasoning_text, _visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + _plain_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _plain_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = full_text), + message = CompletionMessage(**_plain_msg_kwargs), finish_reason = "stop", ) ], ) - api_monitor.set_reply(monitor_id, full_text) + api_monitor.set_reply(monitor_id, _visible_text) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -7790,10 +7925,18 @@ def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: class _ResponsesReasoningExtractor: """Split local markup into Responses reasoning and visible text.""" - def __init__(self, *, parse_think_markers: bool = False) -> None: + def __init__( + self, + *, + parse_think_markers: bool = False, + reasoning_prefilled: bool = False, + ) -> None: self._buffer = "" - self._in_reasoning = False - self._parse_think_markers = parse_think_markers + # ``reasoning_prefilled``: output begins inside an unclosed ```` (Qwen3/GLM prefill), + # so start in reasoning to capture leading text until the first ````. + self._in_reasoning = reasoning_prefilled + # Splitting requires marker parsing; a prefilled open implies it. + self._parse_think_markers = parse_think_markers or reasoning_prefilled def feed( self, @@ -7816,14 +7959,21 @@ class _ResponsesReasoningExtractor: if self._in_reasoning: close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) if close_idx != -1: - reasoning_parts.append(self._buffer[:close_idx]) + reasoning_parts.append( + self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "") + ) self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + # Hold back a trailing partial of either marker: the close (clean chunk-boundary split) + # and a stray open (so a re-emitted ```` isn't leaked into the reasoning drawer). + keep = _responses_marker_holdback( + self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) + ) if keep == len(self._buffer): break - reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + emit = self._buffer[:-keep] if keep else self._buffer + reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, "")) self._buffer = self._buffer[-keep:] if keep else "" break @@ -7860,7 +8010,7 @@ class _ResponsesReasoningExtractor: return "", remaining if self._in_reasoning: self._in_reasoning = False - return remaining, "" + return remaining.replace(_RESPONSES_THINK_OPEN, ""), "" return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") @@ -7869,8 +8019,12 @@ def _extract_responses_reasoning( reasoning_content: Any = None, *, parse_think_markers: bool = False, + reasoning_prefilled: bool = False, ) -> tuple[str, str]: - extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + extractor = _ResponsesReasoningExtractor( + parse_think_markers = parse_think_markers, + reasoning_prefilled = reasoning_prefilled, + ) reasoning, visible = extractor.feed(text, reasoning_content) final_reasoning, final_visible = extractor.finish() return reasoning + final_reasoning, visible + final_visible @@ -9700,7 +9854,7 @@ async def anthropic_messages( # Strip stale tool-call XML from conversation for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml(_msg["content"]).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -9854,7 +10008,7 @@ async def _anthropic_tool_stream( # content event that was purely tool XML doesn't count as text. if etype == "content": event = dict(event) - event["text"] = _TOOL_XML_RE.sub("", event["text"]) + event["text"] = _strip_tool_xml(event["text"]) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -10040,7 +10194,7 @@ async def _anthropic_tool_non_streaming( etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML - clean = _TOOL_XML_RE.sub("", event["text"]) + clean = _strip_tool_xml(event["text"]) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10509,10 +10663,11 @@ async def _anthropic_passthrough_non_streaming( else: text = message.get("content") or "" if text: - # Keep unpromoted bytes when healing is active; legacy stripping is - # only for opted-out or no-client-tool requests. + # Keep unpromoted bytes when healing is active; legacy stripping is only for opted-out + # or no-client-tool requests. _strip_tool_xml also cleans Mistral [TOOL_CALLS] and + # guarded function-XML, not just _TOOL_XML_RE. if not healing_active: - text = _TOOL_XML_RE.sub("", text) + text = _strip_tool_xml(text) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 8df8d37a52..63df86ec17 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -21,7 +21,10 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.tool_call_parser import parse_tool_calls_from_text +from core.inference.tool_call_parser import ( + _gemma_parse_value, + parse_tool_calls_from_text, +) def _args(call: dict) -> dict: @@ -45,6 +48,17 @@ def test_normal_multi_key_arguments_still_split(): assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} +def test_empty_bare_value_becomes_empty_string_not_dropped(): + # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON). + calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"query": "", "unit": "celsius"} + + only = parse_tool_calls_from_text("<|tool_call>call:get{q:}") + assert len(only) == 1, only + assert _args(only[0]) == {"q": ""} + + def test_bare_value_with_timestamps_after_comma_is_kept(): # A comma followed by digits-then-colon (a timestamp/ratio) is value text, # not a new key, so the whole query must be preserved as one argument. @@ -159,3 +173,43 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_gemma_parse_value_always_advances_on_stray_delimiter(): + # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the + # parser, or a looping caller spins forever (DoS). + for delim in (",", "}", "]"): + text = delim + "rest" + value, nxt = _gemma_parse_value(text, 0) + assert nxt > 0, (delim, value, nxt) + + +def test_malformed_gemma_array_does_not_hang(): + # ``[},]`` (stray ``}`` in a list body) hung the buggy parser; the timeout fails + # the regression loudly instead of blocking CI forever. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input" + + +def test_malformed_gemma_mapping_value_does_not_hang(): + # A stray ``}`` where a mapping value is expected must also terminate. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input" diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 05d2a0b80a..8977d6e92a 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -20,7 +20,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.llama_cpp import _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend +from core.inference.llama_cpp import ( + _MAX_REPROMPTS, + _PROVISIONAL_ARGS_MIN_CHARS, + LlamaCppBackend, +) from state import tool_approvals from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision @@ -1036,9 +1040,11 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): """No-tool re-prompt attempts should not concatenate into the UI.""" - streams = [ - [_sse({"content": "I will use render_html now."}), _done()], - [_sse({"content": "Understood. I will use render_html now."}), _done()], + # One initial response plus one stream per re-prompt (count from the shared cap). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + streams += [ + [_sse({"content": "Understood. I will use render_html now."}), _done()] + for _ in range(_MAX_REPROMPTS) ] payloads: list[dict] = [] backend = _make_backend(monkeypatch, streams, payloads) @@ -1073,7 +1079,7 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now."] - assert len(payloads) == 2 + assert len(payloads) == _MAX_REPROMPTS + 1 def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): @@ -1200,6 +1206,66 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc ) +def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch): + # Inline Mistral ``[TOOL_CALLS]`` after a visible preface: the DRAINING flush must use the + # shared parser patterns (the legacy set leaked the marker to clients). + streams = [ + [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "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 = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("[TOOL_CALLS]" not in t for t in content_texts), content_texts + assert any("Let me search." in t for t in content_texts) + + +def test_textual_llama_python_tag_marker_not_leaked(monkeypatch): + # Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form. + streams = [ + [_sse({"content": '<|python_tag|>web_search.call(query="cats")'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "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 = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("<|python_tag|>" not in t for t in content_texts), content_texts + + def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): """Suppression ends once a forced re-prompt actually calls a tool.""" @@ -1738,6 +1804,189 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch): assert calls == [("python", {"code": big_code})] +def _streamed_content(text: str, frag: int = 4) -> list[str]: + """Stream content token-by-token like llama-server; ``frag`` sets the chunk size.""" + chunks = [_sse({"content": text[i : i + frag]}) for i in range(0, len(text), frag)] + chunks.append(_done()) + return chunks + + +def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch): + """A wrapper-less bare-JSON call must be held while incomplete, drained silently, and executed with nothing leaking.""" + + bare_call = '{"name": "web_search", "parameters": {"query": "weather in Sydney"}}' + first_stream = _streamed_content(bare_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "weather in Sydney"})] + assert any( + event.get("type") == "tool_end" and event.get("tool_name") == "web_search" + for event in events + ) + + # The bare JSON never leaked to the user-visible stream. + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('"name"' not in t for t in content_texts), content_texts + assert all("web_search" not in t for t in content_texts), content_texts + # The post-tool synthesis is still streamed. + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypatch): + """Markerless JSON with a non-enabled name is the answer, not a phantom call.""" + + answer = '{"name": "Alice", "parameters": {"age": 30}}' + first_stream = _streamed_content(answer) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give me a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): + """If generation is cut off mid bare-JSON object (no closing brace), the held + fragment must be stripped at stream end rather than dumped to the user.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch): + """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names).""" + + truncated = '{"name": "Alice", "parameters": {"age": 30' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_enabled_name_json_is_still_suppressed(monkeypatch): + """Counterpart guard: a truncated ENABLED-tool bare call (``web_search``) cut off + mid-JSON still must NOT leak -- the gate only spares disabled / non-tool names.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): + """An oversized still-open JSON answer with a non-enabled name streams as content, not a phantom drain.""" + + cap = 16384 + big = "A" * (cap + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)] + first_stream.append(_done()) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "long json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts[:1] + + def _usage_done(usage: dict, finish_reason: str = "stop") -> str: """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the real server reports it on the final chunk of a completion.""" @@ -1813,3 +2062,131 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): metadata = [e for e in events if e.get("type") == "metadata"] assert metadata, "expected a metadata event" assert "prompt_tokens_details" not in metadata[-1]["usage"] + + +def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): + """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" + + cap = 16384 + big = "A" * (cap + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)] + first_stream.append(_done()) + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1] + assert calls and calls[0][0] == "python" + assert len(calls[0][1].get("code", "")) > cap + + +def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch): + """After a bare-JSON call executes, the kept assistant message must not carry the raw call as content.""" + + import copy + + first_stream = [ + _sse({"content": '{"name":"web_search","parameters":{"query":"cats"}}'}), + _done(), + ] + final_stream = [_sse({"content": "Found."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "RESULT") + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert len(payloads) >= 2 + asst = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + +def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): + """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" + + trunc = '{"name":"web_search","parameters":{"query":"weather' + + def _run(auto_heal): + stream = [_sse({"content": trunc}), _done()] + backend = _make_backend(monkeypatch, [stream], []) + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + auto_heal_tool_calls = auto_heal, + ) + ) + contents = "".join(e.get("text", "") for e in events if e.get("type") == "content") + return calls, contents + + calls_off, contents_off = _run(False) + assert calls_off == [], calls_off + assert "web_search" in contents_off, contents_off + + calls_on, contents_on = _run(True) + assert calls_on == [], calls_on + assert "web_search" not in contents_on, contents_on + + +def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): + """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds.""" + # More tool-call streams than the budget: leaked re-prompt slots would run 2+3=5 rounds; + # honouring the budget stops after 2, then a tool-less final-answer pass. + streams = [ + _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search repeatedly"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # Exactly two executed tool rounds, then one final-answer pass. + assert len(calls) == 2, calls + assert len(payloads) == 3, len(payloads) + # The final pass is the budget-exhausted nudge and carries no tools. + assert _tool_names(payloads[2]) == [], _tool_names(payloads[2]) + assert any( + m.get("role") == "user" and "used all available tool calls" in m.get("content", "") + for m in payloads[2]["messages"] + ), payloads[2]["messages"] diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index a7ceb49ed9..ce5688be3e 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -59,6 +59,7 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _ResponsesReasoningExtractor, _SameTaskStreamingResponse, _build_chat_request, _chat_tool_calls_to_responses_output, @@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter: def test_monitor_records_translated_visible_text(self, monkeypatch): import routes.inference as inf_mod + import routes.inference as inf_mod async def fake_chat_completions(chat_req, request): assert request.state.skip_api_monitor is True @@ -1988,6 +1990,122 @@ class TestTranslatedMessagesValidate: ChatMessage(**m.model_dump(exclude_none = True)) +# reasoning_prefilled: Qwen3/GLM enable_thinking templates prefill an unclosed , so generation +# begins inside the think block and emits only the closing ; extractor starts in reasoning. +class TestReasoningPrefilledExtractor: + def test_prefilled_single_feed_splits_lone_close(self): + # T1: reasoning...answer with a prefilled (unseen) open tag. + reasoning, visible = _extract_responses_reasoning( + "plananswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "answer" + + def test_prefilled_never_closed_is_all_reasoning(self): + # T2: truncated mid-thought (no ) -> all reasoning (GGUF parity). + reasoning, visible = _extract_responses_reasoning( + "still thinking with no close", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "still thinking with no close" + assert visible == "" + + def test_prefilled_close_split_across_feeds(self): + # T3: straddles two feed() calls; holdback resolves it. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + r1, v1 = ex.feed("planans") + fr, fv = ex.finish() + assert (r1 + r2 + fr) == "plan" + assert (v1 + v2 + fv) == "ans" + + def test_prefilled_close_split_one_char_per_feed(self): + # T4: every char in its own feed still splits correctly. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + reasoning, visible = "", "" + for ch in "planx": + r, v = ex.feed(ch) + reasoning += r + visible += v + fr, fv = ex.finish() + assert (reasoning + fr) == "plan" + assert (visible + fv) == "x" + + def test_prefilled_empty_generation(self): + # T5: nothing generated. + reasoning, visible = _extract_responses_reasoning( + "", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "" + + def test_prefilled_whitespace_after_close_is_visible(self): + # T6: Qwen commonly emits \n\n before the answer. + reasoning, visible = _extract_responses_reasoning( + "plan\n\nanswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "\n\nanswer" + + def test_prefilled_stray_open_tag_is_suppressed(self): + # T7: a re-emitted literal inside prefilled reasoning is dropped, not leaked. + reasoning, visible = _extract_responses_reasoning( + "abc", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "ab" + assert visible == "c" + assert "" not in reasoning + + def test_prefilled_close_at_start_empty_reasoning(self): + # T8: model closed immediately (empty reasoning) then answered. + reasoning, visible = _extract_responses_reasoning( + "hi", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "hi" + + def test_not_prefilled_lone_close_preserves_current_behavior(self): + # T9: without prefilled, a lone keeps pre-fix behavior (reasoning stays visible, tag dropped). + reasoning, visible = _extract_responses_reasoning( + "reasoningans", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "reasoningans" + + def test_not_prefilled_full_pair_still_splits(self): + # T10: normal explicit .. (GGUF / Harmony) unchanged. + reasoning, visible = _extract_responses_reasoning( + "rv", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "r" + assert visible == "v" + + def test_prefilled_ignored_when_markers_not_parsed(self): + # T11: a non-reasoning model (parse_think_markers False) passes text straight through. + reasoning, visible = _extract_responses_reasoning( + "just an answer", + parse_think_markers = False, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "just an answer" + + # ===================================================================== # Streaming passthrough healing — text-form calls promoted in order # ===================================================================== diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 671af93708..643d64af7a 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -127,9 +127,8 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS], -# which our parser can't read. The route helper must not flip supports_tools=True -# for them, else the UI enables a pill the agentic loop can't honour. +# Llama-3 / Mistral / Gemma 4 tool-call formats are parser-supported, so supports_tools stays True; +# only templates matching none of the known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -161,27 +160,106 @@ MISTRAL_TEMPLATE = """ {%- endfor %} """ +GEMMA4_TEMPLATE = """ +{%- if tools %} + {{- 'Tools available. Emit calls as ' }} + {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" -def test_detect_safetensors_features_llama3_template_suppresses_tools(): - """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + +def test_detect_safetensors_features_llama3_template_keeps_tools_on(): + """Llama-3 emits <|python_tag|>; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) - assert flags["supports_tools"] is False + assert flags["supports_tools"] is True -def test_detect_safetensors_features_mistral_template_suppresses_tools(): - """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" +def test_detect_safetensors_features_mistral_template_keeps_tools_on(): + """Mistral emits [TOOL_CALLS]; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): + """Gemma 4 emits <|tool_call>; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit") + flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE) + assert flags["supports_tools"] is True + + +LLAMA3_2_BARE_JSON_TEMPLATE = """ +{%- if tools %} + {{- 'Given the following functions, respond with JSON for a function call.' }} + {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if 'tool_calls' in message %} + {{- '{"name": "' + message.tool_calls[0].function.name + '", '}} + {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on(): + """Llama-3.2 bare JSON is supported, so the pill stays enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +MINICPM5_ATTRIBUTE_TEMPLATE = """ +{%- if tools %} + {{- 'Available tools. Emit calls as ' }} + {{- 'value' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_attribute_function_form_keeps_tools_on(): + """The attribute form ```` must be whitelisted or the pill is wrongly suppressed.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5") + flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_unknown_format_suppresses_tools(): + """Tools advertised with no known marker must be suppressed.""" + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}<|im_start|>system\n" + "Emit tool calls as JSON-RPC notifications inside the response." + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format") + flags = _detect_safetensors_features(backend, tpl) assert flags["supports_tools"] is False def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): - """Sanity check: gate only suppresses non-Qwen formats.""" + """Sanity check: Qwen marker still flips supports_tools.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") @@ -454,3 +532,130 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): assert flags["supports_tools"] is True assert flags["supports_reasoning"] is True assert flags["supports_preserve_thinking"] is True + + +# Templates advertising tools whose ``{"name":`` example is pretty-printed or JSON-escaped. +_WHITESPACE_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + "To call a tool, output JSON of the form:\n" + '{ "name" : "function_name", "parameters": { } }\n' + "{%- endif %}\n" + "{{ messages }}" +) +_ESCAPED_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + 'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n' + "{%- endif %}\n" + "{{ messages }}" +) +_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = ( + "{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}" +) + + +def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json(): + # Pretty-printed bare-JSON (``{ "name" :``) keeps supports_tools: parser accepts the whitespace. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json(): + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): + # Negative control: tools advertised but no parser-recognised emission form -> pill dropped. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): + # The {"function":...} bare-JSON alias keeps supports_tools, mirroring {"name":...}. + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}\n" + 'Respond with {"function": "fn", "parameters": {}}\n' + "{%- endif %}\n" + "{{ messages }}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# _sf_reasoning_prefill_mode gates the prefilled- extractor for enable_thinking models. +class TestSafetensorsReasoningPrefillGate: + # Qwen3-style template with the standard / markers. + _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" + # gemma-style bespoke reasoning channel -- no standard markers. + _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + + def _features(self, **over): + base = { + "supports_reasoning": True, + "reasoning_always_on": False, + "reasoning_style": "enable_thinking", + } + base.update(over) + return base + + def test_g1_enable_thinking_true(self): + # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + + def test_g2_enable_thinking_none_defaults_on(self): + # G2: default request (None) -> prefilled (Qwen3/GLM templates default on). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + + def test_g3_enable_thinking_false(self): + # G3: thinking explicitly off -> not prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + + def test_g4_gpt_oss_reasoning_effort_excluded(self): + # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "reasoning_effort") + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g5_enable_thinking_effort_included(self): + # G5: GLM-style enable_thinking_effort also prefills. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "enable_thinking_effort") + assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True + + def test_g6_non_reasoning_model(self): + # G6: no reasoning capability -> never prefilled. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(supports_reasoning = False, reasoning_style = None) + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g7_reasoning_always_on(self): + # G7: hardcoded- template -> prefilled regardless of the flag. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True + + def test_g8_gemma_bespoke_channel_excluded(self): + # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled (else the + # whole answer is swallowed as reasoning). Regression guard. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + + def test_g9_missing_template_not_prefilled(self): + # G9: no template available -> conservative (not prefilled). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, None) is False diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py new file mode 100644 index 0000000000..9158d1ad5e --- /dev/null +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Safetensors/MLX reasoning-block parity with GGUF. + +enable_thinking templates prefill an unclosed ````, so the stream must split the leading +text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy +of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import ( + _ResponsesReasoningExtractor, + _sf_reasoning_prefill_mode, + _strip_tool_xml_for_display, +) + + +def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: + """Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end.""" + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + reasoning_deltas: list[str] = [] + visible_deltas: list[str] = [] + monitor: list[str] = [] + tool_starts: list[dict] = [] + order: list[str] = [] # "reasoning" | "visible" | "tool_start" sequence + + def _flush(): + fr, fv = extractor.finish() + if fr: + reasoning_deltas.append(fr) + order.append("reasoning") + if fv: + visible_deltas.append(fv) + monitor.append(fv) + order.append("visible") + + for event in events: + etype = event["type"] + if etype == "status": + if not event["text"]: + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + tool_starts.append(event) + order.append("tool_start") + continue + clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True) + new_text = clean[len(prev_text) :] + prev_text = clean + if not new_text: + continue + r, v = extractor.feed(new_text) + if r: + reasoning_deltas.append(r) + order.append("reasoning") + if v: + visible_deltas.append(v) + monitor.append(v) + order.append("visible") + _flush() + return { + "reasoning": "".join(reasoning_deltas), + "visible": "".join(visible_deltas), + "monitor": "".join(monitor), + "tool_starts": tool_starts, + "order": order, + } + + +def test_s1_plain_stream_splits_prefilled_reasoning(): + # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only. + events = [ + {"type": "content", "text": "Let me compute 17*23"}, + {"type": "content", "text": "Let me compute 17*23 = 391The answer is 391."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "Let me compute 17*23 = 391" + assert out["visible"] == "The answer is 391." + assert out["monitor"] == "The answer is 391." + assert "" not in out["reasoning"] and "" not in out["visible"] + + +def test_s2_reasoning_flushed_before_tool_start(): + # S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start. + events = [ + {"type": "content", "text": "I should search"}, + {"type": "content", "text": "I should search Sydney weather"}, + {"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "Found itSydney is 21C today."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + # Both turns' reasoning surfaced, answer only from turn 2. + assert "I should search Sydney weather" in out["reasoning"] + assert "Found it" in out["reasoning"] + assert out["visible"] == "Sydney is 21C today." + assert out["monitor"] == "Sydney is 21C today." + # Ordering: the pre-tool reasoning is emitted before the tool_start. + assert out["order"].index("reasoning") < out["order"].index("tool_start") + + +def test_s3_extractor_resets_each_turn(): + # S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each). + events = [ + {"type": "content", "text": "turn1 thoughtspartial"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "turn2 thoughtsfinal answer"}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "turn1 thoughtsturn2 thoughts" + assert out["visible"] == "partialfinal answer" + + +def test_s4_harmony_full_tags_normal_mode(): + # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False). + events = [{"type": "content", "text": "reasoning herevisible answer"}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "reasoning here" + assert out["visible"] == "visible answer" + + +def test_s5_thinking_off_no_reasoning_deltas(): + # S5: thinking disabled -> not prefilled, no , all content is visible. + events = [{"type": "content", "text": "Just the plain answer, no thinking."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "" + assert out["visible"] == "Just the plain answer, no thinking." + assert out["monitor"] == "Just the plain answer, no thinking." + + +_THINK_TPL = "...{% if enable_thinking %}{% endif %}......" + + +def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): + # GLM-5.2 enable_thinking_effort + reasoning_effort="none" disables thinking like + # enable_thinking=False, so prefilled must be OFF (else the answer is swallowed into reasoning). + feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False + # Thinking on (effort level or default) still prefills. + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True + # An explicit enable_thinking=False also disables (unchanged). + assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False + # reasoning_always_on wins regardless of reasoning_effort. + always = {**feats, "reasoning_always_on": True} + assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True + # Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected. + plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True + + # End-to-end: with prefilled=False, a plain no- answer stays visible. + events = [{"type": "content", "text": "The capital of France is Paris."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["visible"] == "The capital of France is Paris." + assert out["reasoning"] == "" + # The buggy prefilled=True path is what swallowed the whole answer (guard the delta). + swallowed = _replay_sf_reasoning_stream(events, prefilled = True) + assert swallowed["visible"] == "" + assert swallowed["reasoning"] == "The capital of France is Paris." diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 3f2d49f0dd..984d5f8ae9 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -138,6 +138,20 @@ class TestParser: assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + # Only the wrapping newline is trimmed, so code indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_function_signal_inside_parameter_is_literal(self): text = ( "" @@ -189,6 +203,13 @@ class TestParser: text = 'before <|tool_call>call:terminal{command:"ls"} after' assert strip_tool_markup(text) == "before after" + def test_strip_named_mistral_call_consumes_trailing_eos(self): + # The named [TOOL_CALLS]name{json} shape must eat the optional trailing . + text = '[TOOL_CALLS]web_search{"query":"cats"}' + assert strip_tool_markup(text) == "" + text = '[TOOL_CALLS]web_search{"query":"cats"} and then' + assert strip_tool_markup(text) == " and then" + def test_strip_markup_unclosed_final(self): text = "before {partial" # final=True drops the trailing run. @@ -214,6 +235,376 @@ class TestParser: == "before " ) + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): + # A literal in a value is data: the strip closes at the REAL , keeping prose. + raw = ( + "pref " + 'print("") tail' + ) + assert strip_tool_markup_streaming(raw) == "pref tail" + # Streaming and final strip agree on the visible text (final also trims). + assert strip_tool_markup_streaming(raw) == strip_tool_markup(raw, final = True) + + def test_streaming_strip_drops_leading_magistral_reasoning(self): + # Magistral reasoning is a leading [THINK]...[/THINK] block; the streaming strip must drop it. + closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." + assert strip_tool_markup_streaming(closed) == "The answer is 4." + assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) + # Unclosed mid-stream reasoning is held; cleaned text grows only after [/THINK]. + assert strip_tool_markup_streaming("[THINK]still thinking") == "" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The") == "The" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The answer") == "The answer" + # A non-leading [THINK] is ordinary prose, left untouched. + assert strip_tool_markup_streaming("hi [THINK] later") == "hi [THINK] later" + + +class TestParserMultiFormat: + """Shared-parser coverage: every family's emission maps to the same OpenAI shape.""" + + # Llama-3 + + def test_llama3_python_tag_dot_call(self): + # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). + import json + + text = '<|python_tag|>brave_search.call(query="weather in Tokyo")' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "brave_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in Tokyo"} + + def test_llama3_python_tag_dot_call_multi_arg(self): + import json + + text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius", "days": 5} + + def test_llama3_python_tag_json_form(self): + import json + + text = '<|python_tag|>{"name":"web_search","parameters":{"query":"hi","n":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "hi", "n": 5} + + def test_llama3_python_tag_json_form_with_eom(self): + # Llama-3 emits <|eom_id|> after the JSON; must not break parsing. + import json + + text = '<|python_tag|>{"name":"python","parameters":{"code":"print(2+2)"}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"code": "print(2+2)"} + + def test_llama3_strip_markup_final(self): + text = '<|python_tag|>brave_search.call(query="x")' + assert strip_tool_markup(text, final = True) == "" + + # Llama-3.2 bare JSON ``custom_tools`` + + def test_llama3_2_bare_json_parameters(self): + # Llama-3.2-Instruct emits bare JSON directly as content, no <|python_tag|> prefix. + import json + + text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "Tokyo weather"} + + def test_llama3_2_bare_json_arguments_key(self): + import json + + text = '{"name":"add","arguments":{"a":1,"b":2}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"a": 1, "b": 2} + + def test_llama3_2_bare_json_multi_call(self): + # Llama-3 may chain calls with "; " per training template. + text = '{"name":"a","parameters":{}}; {"name":"b","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_llama3_2_bare_json_with_eom_sentinel(self): + text = '{"name":"x","parameters":{"y":1}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_leading_sentinel_skipped(self): + # Sometimes prior <|eot_id|> leaks into the next turn. + text = '<|eot_id|>{"name":"x","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_plain_prose_does_not_fire(self): + # Defensive: must NOT fire on plain assistant prose. + text = "Hello world, how are you today?" + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): + # Defensive: JSON embedded in prose must NOT fire (content must START with `{`). + text = 'The tool result was: {"name":"foo"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_name_does_not_fire(self): + text = '{"result":"ok","data":[1,2,3]}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_args_does_not_fire(self): + text = '{"name":"x"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_args_not_dict_does_not_fire(self): + text = '{"name":"x","parameters":42}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_parameters_does_not_fire(self): + # Llama-3 spec: parameters must be a dict; a string value must NOT trigger. + text = '{"name":"foo","parameters":"this is a sentence"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): + # OpenAI arguments may be a JSON-string of a dict, but a plain non-JSON string must not pass. + text = '{"name":"foo","arguments":"not json"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_json_dict_fires(self): + # OpenAI shape: arguments is a JSON-encoded string of a dict. + text = '{"name":"foo","arguments":"{\\"q\\":\\"x\\"}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + # arguments stays as the original JSON-string. + assert result[0]["function"]["arguments"] == '{"q":"x"}' + + def test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire(self): + # JSON-string that parses to a list / scalar / null must NOT fire. + for bad in ( + '{"name":"foo","arguments":"[1,2,3]"}', + '{"name":"foo","arguments":"\\"plain\\""}', + '{"name":"foo","arguments":"null"}', + '{"name":"foo","arguments":"42"}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # Mistral pre-v11 + + def test_mistral_pre_v11_array(self): + import json + + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"query":"hello"},"id":"abc"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + # Mistral provides its own id; preserve it. + assert result[0]["id"] == "abc" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "hello"} + + def test_mistral_array_parameters_key_alias(self): + import json + + # Array object keyed on parameters (not arguments) must keep its payload. + text = '[TOOL_CALLS] [{"name":"get_weather","parameters":{"city":"Paris"}}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Paris"} + + def test_mistral_pre_v11_array_multi(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"x":1},"id":"id1"},' + '{"name":"b","arguments":{"y":2},"id":"id2"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_mistral_pre_v11_unclosed_array(self): + # Closing ] truncated: parser must heal off individual objects. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Mistral v11+ + + def test_mistral_v11_single(self): + # Magistral / Mistral Small 3.1: bare name{json} after trigger. + import json + + text = '[TOOL_CALLS]add{"a":3.5,"b":4}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} + + def test_mistral_v11_parallel(self): + # v11+ parallel: [TOOL_CALLS]a{...}[TOOL_CALLS]b{...}. + text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "add" + assert result[1]["function"]["name"] == "sub" + + def test_mistral_v11_with_args_marker(self): + # Ministral / Mistral Large 3: [TOOL_CALLS]name[ARGS]{json}. + import json + + text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 1, "b": 2} + + def test_mistral_strip_markup_v11(self): + text = '[TOOL_CALLS]add{"a":1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_call_id_form(self): + # Mistral Small 3.2: the [CALL_ID] segment must be skipped, not treated as a stop (llama.cpp test-chat.cpp:4785). + import json + + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_mistral_call_id_form_parallel(self): + text = ( + '[TOOL_CALLS]special_function[CALL_ID]000000001[ARGS]{"arg1": 1}' + "[TOOL_CALLS]special_function_with_opt[CALL_ID]000000002" + '[ARGS]{"arg1": 1, "arg2": 2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "special_function" + assert result[1]["function"]["name"] == "special_function_with_opt" + + def test_mistral_call_id_form_stripped(self): + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_think_reasoning_ignored(self): + # A [TOOL_CALLS] inside [THINK]...[/THINK] is reasoning; only the call after [/THINK] counts (llama.cpp test-chat.cpp:2285). + import json + + text = ( + '[THINK]Let me think about [TOOL_CALLS]fake[ARGS]{"x":1} ' + 'and more[/THINK][TOOL_CALLS]real_fn[ARGS]{"y":2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "real_fn" + assert json.loads(result[0]["function"]["arguments"]) == {"y": 2} + + def test_mistral_think_reasoning_no_real_call(self): + # Reasoning that mentions a call but emits none after [/THINK] yields no calls. + text = '[THINK]I might call [TOOL_CALLS]fake[ARGS]{"x":1}[/THINK]Done.' + assert parse_tool_calls_from_text(text) == [] + + def test_mistral_think_literal_in_argument_preserved(self): + # A literal [THINK] inside a real tool argument must not be stripped or corrupt the parse. + import json + + text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"q": "explain the [THINK] token"} + + # Gemma 4 + + def test_gemma4_simple_call(self): + import json + + text = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>Tokyo<|"|>,units:<|"|>celsius<|"|>}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius"} + + def test_gemma4_with_primitives(self): + import json + + text = ( + "<|tool_call>call:set_pref{" + "enabled:true,attempts:5,threshold:1.5,nickname:null}" + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} + + def test_gemma4_nested_args(self): + # Gemma 4 nests dicts / lists with bare keys and <|"|> strings. + import json + + text = ( + "<|tool_call>call:search{" + 'query:<|"|>foo<|"|>,filters:{site:<|"|>example.com<|"|>,recent:true},' + 'tags:[<|"|>a<|"|>,<|"|>b<|"|>]}' + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "foo" + assert args["filters"] == {"site": "example.com", "recent": True} + assert args["tags"] == ["a", "b"] + + def test_gemma4_multi_call(self): + text = "<|tool_call>call:a{x:1}<|tool_call>call:b{y:2}" + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_gemma4_unclosed_does_not_raise(self): + # Truncated mid-stream; must not raise. + text = '<|tool_call>call:foo{x:<|"|>bar<|"|>' + result = parse_tool_calls_from_text(text) + assert isinstance(result, list) + + def test_gemma4_strip_markup_final(self): + text = "<|tool_call>call:foo{x:1}" + assert strip_tool_markup(text, final = True) == "" + + # Cross-format sentinels + + def test_all_markers_in_tool_xml_signals(self): + # Streaming buffer wakes up on every emission marker. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_has_tool_signal_for_all_formats(self): + assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') + assert has_tool_signal('[TOOL_CALLS] [{"name":"x"}]') + assert has_tool_signal('[TOOL_CALLS]add{"a":1}') + assert has_tool_signal("<|tool_call>call:foo{}") + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop @@ -347,6 +738,130 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): + # A late unclosed heals only with Auto-Heal on; off, it must not execute. + prose = "Sure, let me look that up for you right now. " + incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + + loop_off, exec_off = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = False, + max_tool_iterations = 3, + ) + events_off = _collect_events(loop_off) + assert exec_off.calls == [], "disabled Auto-Heal must not execute a healed incomplete call" + assert not [e for e in events_off if e.get("type") == "tool_start"] + + loop_on, exec_on = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = True, + max_tool_iterations = 3, + ) + _collect_events(loop_on) + assert exec_on.calls == [("web_search", {"query": "weather in Sydney"})], exec_on.calls + + +def test_bare_json_tool_call_is_not_streamed_as_content(): + # Llama-3.2 bare form carries no XML signal: BUFFER until the object closes, never leak the JSON. + bare = '{"name":"web_search","parameters":{"query":"cats"}}' + loop, exec_fn = _make_loop( + turns = [[bare], ["Here are the results."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + assert any("Here are the results." in t for t in contents) + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): + # Markerless JSON whose "name" is not an enabled tool must be shown, not dropped. + answer = '{"name":"Alice","parameters":{"age":30}}' + loop, exec_fn = _make_loop(turns = [[answer]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): + # Same as above but the bare object arrives split mid-key, held across chunks until it balances. + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_', 'search","parameters":{"query":"cats"}}'], + ["Done."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + + +def test_leading_json_answer_is_not_dropped(): + # A leading {...} that is NOT a call must still surface; the hold only delays it. + obj = '{"answer": 42, "note": "done"}' + loop, exec_fn = _make_loop( + turns = [[obj]], + exec_results = [], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert any('"answer"' in t for t in contents), contents + + +def _reprompt_loop(*, auto_heal_tool_calls): + """Drive one restricted tool with an intent-only first turn to exercise the nudge; returns conversations and events.""" + captured: list[list] = [] + + def fake_single_turn(messages, active_tools = None): + captured.append(list(messages)) + if len(captured) == 1: + yield "I'll search for that now." # forward-looking intent, no call + else: + yield "Final answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "find X"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + execute_tool = exec_fn, + auto_heal_tool_calls = auto_heal_tool_calls, + max_tool_iterations = 3, + ) + ) + return captured, events + + +def test_reprompt_names_only_active_tools_not_hardcoded(): + # The nudge must name the tools actually enabled, not hardcoded web_search/python. + captured, _events = _reprompt_loop(auto_heal_tool_calls = True) + assert len(captured) >= 2, "intent prose should have triggered a re-prompt turn" + reprompt = captured[1][-1] + assert reprompt["role"] == "user" + assert "search_knowledge_base" in reprompt["content"] + assert "web_search" not in reprompt["content"] + assert "python" not in reprompt["content"] + + +def test_reprompt_suppressed_when_auto_heal_disabled(): + # With Auto-Heal off the nudge stays silent for GGUF parity, so only the initial generation runs. + captured, events = _reprompt_loop(auto_heal_tool_calls = False) + assert len(captured) == 1, captured + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("search for that" in t for t in contents) + + class TestLoopBasic: def test_plain_answer(self): # No tool XML; loop should yield content then status="". @@ -406,6 +921,85 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_llama3_python_tag_form(self): + # The loop must recognise Llama-3's <|python_tag|> marker, drain the turn, and execute the call. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|python_tag|>web_search.call(", + 'query="weather in Tokyo"', + ")", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in Tokyo"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_llama3_bare_json_form_fires_tool(self): + # Llama-3.1/3.2 bare-JSON calls carry no XML signal; the safety-net parse must still fire + # the tool. Regression for the has_tool_signal gate that dropped these. + loop, exec_fn = _make_loop( + turns = [ + ['{"name": "web_search", "parameters": {"query": "weather in SF"}}'], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 18C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in SF"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_mistral_pre_v11_form(self): + # Pre-v11 Mistral emission: [TOOL_CALLS] [{...}]. + loop, exec_fn = _make_loop( + turns = [ + [ + '[TOOL_CALLS] [{"name":"web_search",', + '"arguments":{"query":"hi"},"id":"abc"}]', + ], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + # Mistral-provided ids must propagate to tool_start events. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "abc" + + def test_mistral_v11_form(self): + # v11+ Mistral emission: bare name{json} after the trigger. + loop, exec_fn = _make_loop( + turns = [ + ['[TOOL_CALLS]web_search{"query":"hi"}'], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + + def test_gemma4_form(self): + # Gemma 4 emission: <|tool_call>call:NAME{...}. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_call>call:web_search{", + 'query:<|"|>weather<|"|>', + "}", + ], + ["sunny"], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -765,6 +1359,55 @@ class TestLoopBehaviour: assert len(duplicate_nudges) == 1 assert captured_tool_names[2] == ["web_search", "python"] + def test_duplicate_noop_does_not_consume_budget_at_small_cap(self): + # A duplicate no-op turn must NOT spend the tool budget: only turns that execute a tool + # count (GGUF parity), so a distinct call can still follow at max_tool_iterations=2. + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_tool_names.append( + [ + tool["function"]["name"] + for tool in (active_tools or []) + if tool.get("function", {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result", "python-result"]) + _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + + # Both distinct tools execute; the repeated call in between did not cost a slot. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + # The turn after the duplicate still offered tools (budget not yet spent). + assert captured_tool_names[2] == ["web_search", "python"] + def test_repeated_duplicate_noop_transitions_to_final_attempt(self): captured_tool_names: list[list[str]] = [] turns = iter( @@ -953,6 +1596,234 @@ class TestLoopBehaviour: assert "boom" in tool_end["result"] +class TestLoopRePrompt: + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``_MAX_REPROMPTS`` extra slots.""" + + def test_intent_signal_triggers_reprompt(self): + # Turn 1: intent signal, no tool call. + # Turn 2 (re-prompt): proper tool call -> executes. + # Turn 3: final answer. + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + [ + '{"name":"web_search","arguments":' + '{"query":"sky color"}}' + ], + ["The sky is blue."], + ], + exec_results = ["Blue (Rayleigh scattering)"], + ) + events = _collect_events(loop) + # web_search must have been called once (after the re-prompt). + assert exec_fn.calls == [("web_search", {"query": "sky color"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "blue" in contents[-1]["text"].lower() + + def test_intent_signal_without_tools_does_not_reprompt(self): + # Same intent signal but no tools enabled -- must NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["Let me think about that for a moment."]], + exec_results = [], + ) + # _make_loop hard-codes three tools; rebuild without tools. + from core.inference.safetensors_agentic import run_safetensors_tool_loop + + def _gen(_messages): + yield "Let me think about that for a moment." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + ) + ) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and "think" in contents[-1]["text"].lower() + + def test_direct_answer_does_not_trigger_reprompt(self): + # Plain answer with no intent words: do NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["4"]], + exec_results = [], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and contents[-1]["text"].strip() == "4" + + def test_max_reprompts_capped_at_three(self): + # Model keeps stalling with intent -- after 3 re-prompts the loop must give up. + turns = [["Let me search for that."]] * 6 # well over the cap + loop, exec_fn = _make_loop( + turns = turns, + exec_results = [], + ) + events = _collect_events(loop, max_events = 500) + # No tool ever ran, but the loop terminated cleanly. + assert exec_fn.calls == [] + statuses = [e for e in events if e["type"] == "status"] + assert statuses and statuses[-1]["text"] == "" + + def test_short_intent_below_buffer_threshold_triggers_reprompt(self): + # Short emission that never exits BUFFERING must still trigger the intent re-prompt. + loop, exec_fn = _make_loop( + turns = [ + ["Let me check."], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_reprompt_does_not_consume_tool_budget(self): + # max_tool_iterations=1: the re-prompt must not eat the slot, so the real call still runs. + loop, exec_fn = _make_loop( + turns = [ + # 1. Intent stall (re-prompt 1/3). + ["Let me search for that."], + # 2. Real tool call (uses the budget slot). + ['{"name":"web_search","arguments":{"query":"weather"}}'], + # 3. Budget exhausted -> nudged final answer. + ["Final: it is sunny"], + ], + exec_results = ["sunny"], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + +class TestLoopCanonicalHealKey: + """Per-tool canonical heal key (``code``/``command``/``query``), mirroring GGUF.""" + + def test_python_bare_string_heals_to_code(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"python","arguments":"print(1)"}' ""], + ["done"], + ], + exec_results = ["1\n"], + ) + events = _collect_events(loop) + # The bare string must heal to {"code": ...}, not {"query": ...}, so the python sandbox runs it. + assert exec_fn.calls == [("python", {"code": "print(1)"})] + + def test_terminal_bare_string_heals_to_command(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"terminal","arguments":"ls -la"}' ""], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("terminal", {"command": "ls -la"})] + + def test_unknown_tool_bare_string_heals_to_query(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":"hello"}' ""], + ["ok"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hello"})] + + +class TestGGUFSafetensorsHealingParity: + """Pin GGUF vs safetensors/MLX loop parity so a regression on either side breaks CI.""" + + def test_gguf_imports_shared_signal_markers(self): + # The GGUF BUFFERING machine must wake on every shared emission marker, else calls slip past as prose. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "_SHARED_TOOL_XML_SIGNALS" in src, ( + "GGUF agentic loop must reuse the shared TOOL_XML_SIGNALS " + "tuple so it wakes on all five emission formats" + ) + + def test_gguf_uses_shared_strip_helper(self): + # The GGUF stream-cleanup must delegate to the shared strip_tool_markup for every family. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert ( + "_shared_strip_tool_markup" in src + ), "GGUF stream cleanup must delegate to the shared strip_tool_markup helper" + + def test_gguf_uses_canonical_heal_keys(self): + # GGUF and safetensors heal a bare-string argument to the same canonical key via the shared coerce_tool_arguments. + from core.inference.tool_loop_controller import ( + _CANONICAL_HEAL_ARG, + coerce_tool_arguments, + ) + + assert _CANONICAL_HEAL_ARG["python"] == "code" + assert _CANONICAL_HEAL_ARG["terminal"] == "command" + assert coerce_tool_arguments("print(1)", heal = True, tool_name = "python").arguments == { + "code": "print(1)" + } + assert coerce_tool_arguments("ls -la", heal = True, tool_name = "terminal").arguments == { + "command": "ls -la" + } + assert coerce_tool_arguments("weather", heal = True, tool_name = "web_search").arguments == { + "query": "weather" + } + + def test_intent_regex_matches_same_phrases_as_gguf(self): + # The intent re-prompt regex must match the SAME phrases on both backends. + from core.inference.llama_cpp import _INTENT_SIGNAL as gguf_re + from core.inference.safetensors_agentic import ( + _INTENT_SIGNAL as sf_re, + ) + + for phrase in ( + "I'll search for that", + "I will look it up", + "Let me check", + "I am going to call the tool", + "First, I will explore", + "Here's my plan", + "Now I need to call web_search", + ): + assert gguf_re.search(phrase), f"GGUF missed {phrase!r}" + assert sf_re.search(phrase), f"safetensors missed {phrase!r}" + + for plain in ( + "4", + "Hello!", + "The sky is blue.", + "I can help with that.", + "I should mention", + "Let's go.", + # Negated intent is a refusal, not a plan: neither backend may re-prompt on it. + "I will not search the web for that.", + "I'll never call that tool.", + ): + assert not gguf_re.search(plain), f"GGUF wrongly fired on {plain!r}" + assert not sf_re.search(plain), f"safetensors wrongly fired on {plain!r}" + + def test_max_reprompts_equal_on_both_backends(self): + from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap + from core.inference.safetensors_agentic import _MAX_REPROMPTS as sf_cap + assert gguf_cap == sf_cap == 3 + + class TestLoopControl: def test_cancel_event_breaks_loop(self): cancel = threading.Event() @@ -1407,5 +2278,358 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# Routes-level python_tag strip (multi-line; stop on next sentinel) +class TestRoutesPythonTagStrip: + """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" + + def _strip(self, text: str) -> str: + # Import inside the test so a routes-module import error doesn't fail collection. + from routes.inference import _strip_tool_xml + return _strip_tool_xml(text) + + def test_single_line_python_tag_stripped(self): + # Floor: the original 5620 single-line behaviour still works. + text = '<|python_tag|>brave_search.call(query="weather")' + assert self._strip(text) == "" + + def test_python_tag_with_less_than_in_code(self): + # 5615 regression: a literal < inside code must NOT terminate the strip early. + text = '<|python_tag|>python.call(code="if x < 10: pass")' + assert self._strip(text) == "" + + def test_python_tag_multiline_code_stripped(self): + # 5620 round-1 regression: multi-line code's second line leaked. + text = '<|python_tag|>python.call(code="line1\nline2\nline3")' + assert self._strip(text) == "" + + def test_python_tag_multiline_with_less_than(self): + # Combined: multi-line code AND literal < in code. + text = ( + '<|python_tag|>python.call(code="for i in range(10):\n' + " if i < 5:\n" + ' print(i)")' + ) + assert self._strip(text) == "" + + def test_python_tag_stops_at_eom_sentinel(self): + # Strip stops at the next Llama-3 <| sentinel so trailing assistant content survives. + text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" + assert self._strip(text) == "<|eom_id|>final answer text" + + def test_python_tag_stops_at_eot_sentinel(self): + text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after" + assert self._strip(text) == "<|eot_id|>after" + + def test_python_tag_json_form_multiline_stripped(self): + # The JSON form of python_tag with newlines inside string args. + text = '<|python_tag|>{"name":"python","parameters":{"code":"a = 1\nb = 2\nprint(a+b)"}}' + assert self._strip(text) == "" + + def test_python_tag_with_eom_then_trailing_python_tag(self): + # Two python_tag emissions back-to-back across a sentinel: both strip independently. + text = ( + '<|python_tag|>brave_search.call(query="a")' + "<|eom_id|>" + '<|python_tag|>python.call(code="x=1")' + ) + # <|eom_id|> between the two strips remains; both python_tag blocks are consumed. + assert self._strip(text) == "<|eom_id|>" + + +# Robustness fixes uncovered while validating against vLLM / sglang. +class TestParserRobustness: + def test_tool_call_json_accepts_parameters_key(self): + # Hermes wrapper using parameters instead of arguments; this path now accepts both keys. + import json + + text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "search" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "ramen"} + + def test_function_xml_attribute_form(self): + # MiniCPM-5 / MiniMax-M2 attribute syntax: v. + import json + + text = '' 'Tokyo' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_xml_attribute_form_multi_param(self): + import json + + text = ( + '' + 'Tokyo' + 'celsius' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"city": "Tokyo", "unit": "celsius"} + + def test_function_xml_legacy_equals_form_still_works(self): + # Regression guard: the old v syntax must keep parsing after the regex broadening. + import json + + text = "Tokyo" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_form_has_tool_signal(self): + # The standalone form must flip the streaming buffer, else the call is dropped. + assert has_tool_signal('') is True + + def test_function_attribute_form_strip_markup(self): + # The attribute form must also be stripped from displayed text, like . + text = 'result X' + assert strip_tool_markup(text, final = True) == "result" + + def test_llama3_chat_template_round_trip(self): + # Llama-3.x prefixes assistant turns with <|start_header_id|>...<|end_header_id|>; the + # sentinel-strip must reach past the role label to the JSON body, else history calls drop. + import json + + text = ( + "<|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "get_weather", "parameters": {"city": "Tokyo"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_llama3_round_trip_all_roles(self): + # Same logic must work for every role the chat template inserts. + import json + for role in ("assistant", "user", "system", "tool", "ipython"): + text = ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n" + '{"name": "f", "parameters": {"x": 1}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"failed for role={role}" + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_llama3_round_trip_with_eot_prefix(self): + # Prior turn closes with <|eot_id|>, then the new header opens; both sentinels + role must be consumed. + import json + + text = ( + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "f", "parameters": {}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "f" + + def test_function_xml_followed_by_prose(self): + # Body must terminate at even without a wrapper, else prose leaks into the value. + import json + + text = ( + "" + "Tokyo" + "\n\nHere is what I found." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_xml_followed_by_prose(self): + # Same expectation for the MiniCPM-5 attribute form. + import json + + text = ( + '' + 'Tokyo' + "\n\nLet me know if you need anything else." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +def test_truncated_bare_json_at_eof_is_not_leaked(): + # Stream ends mid bare-JSON: the held fragment must be dropped at EOF, not flushed as content. + loop, _exec = _make_loop( + turns = [['{"name":"web_search","parameters":{"query":"weather in S']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t for t in contents), contents + + +def test_oversized_bare_json_call_is_not_leaked_and_executes(): + # A bare-JSON call exceeding _MAX_BARE_JSON_BUFFER must DRAIN, not stream the prefix, and still execute. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, exec_fn = _make_loop(turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(t.lstrip().startswith('{"name') for t in contents), contents[:1] + assert exec_fn.calls and exec_fn.calls[0][0] == "python" + assert len(exec_fn.calls[0][1].get("code", "")) > _MAX_BARE_JSON_BUFFER + + +def test_oversized_plain_json_answer_still_streams(): + # A giant plain JSON answer (no "name" key) is NOT a call and must still stream. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"result":"' + big + '"}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, _exec = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents + + +def test_oversized_disabled_name_json_answer_still_streams(): + # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream, not drain. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + chunks = [answer[i : i + 2000] for i in range(0, len(answer), 2000)] + loop, exec_fn = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents[:80] + + +def test_truncated_disabled_name_json_is_shown_at_eof(): + # A truncated JSON answer whose name is not an enabled tool must be shown at EOF. + truncated = '{"name":"Alice","parameters":{"age":' + loop, exec_fn = _make_loop(turns = [[truncated]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_truncated_plain_json_with_nested_enabled_name_is_visible(): + # A truncated answer with only a NESTED "name" must be shown: the gate uses the TOP-LEVEL name. + loop, exec_fn = _make_loop( + turns = [['{"result":{"name":"web_search","age":']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents and "web_search" in contents, contents + + +def test_bare_json_call_not_replayed_in_next_turn_content(): + # After a bare-JSON call executes, the next-turn assistant content must not contain the raw call. + captured: list[list[dict]] = [] + exec_fn = FakeExecuteTool(["RESULT"]) + + def st(messages, active_tools = None): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"web_search","parameters":{"query":"cats"}}' + else: + yield "Found." + + _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + assert len(captured) >= 2, captured + asst = [m for m in captured[1] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): + # With Auto-Heal OFF a truncated enabled-name bare-JSON fragment stays visible; with it ON, suppressed. + trunc = '{"name":"web_search","parameters":{"query":"weather' + off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) + events_off = _collect_events(off) + assert exec_off.calls == [], exec_off.calls + contents_off = "".join(e["text"] for e in events_off if e["type"] == "content") + assert "web_search" in contents_off, contents_off + + on, exec_on = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True) + events_on = _collect_events(on) + assert exec_on.calls == [], exec_on.calls + contents_on = "".join(e["text"] for e in events_on if e["type"] == "content") + assert "web_search" not in contents_on, contents_on + + +def test_looks_like_enabled_bare_json_accepts_function_alias(): + # The buffering gate must recognise the "function" bare-JSON alias, so it is buffered, not streamed. + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + + enabled = {"web_search"} + assert _looks_like_enabled_bare_json( + '{"function":"web_search","parameters":{"q":"x"}}', enabled + ) + # A non-tool "function" value is an ordinary JSON answer -> not gated. + assert not _looks_like_enabled_bare_json('{"function":"Alice","parameters":{}}', enabled) + + +class TestFalseAlarmMarkerProse: + def test_leading_marker_prose_streams_intact(self): + # An answer starting with a literal marker is a false alarm: the full prose must reach the client. + text = "[TOOL_CALLS] is the Mistral tool marker. More prose after." + loop, exec_fn = _make_loop(turns = [[text]]) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert texts and texts[-1] == text + + def test_chained_bare_json_calls_not_replayed_in_history(self): + # Both chained calls execute; the next-turn history must not contain the second call's raw JSON. + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + convs = [] + turn_iter = iter([[chained], ["Final answer."]]) + + def gen(messages, active_tools = None): + convs.append([dict(m) for m in messages]) + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["r1", "r2"]) + loop = run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search", "python"] + assistant = next(m for m in convs[1] if m["role"] == "assistant") + assert '"python"' not in (assistant.get("content") or "") diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 39fdd151be..7664126d91 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -102,6 +102,22 @@ class TestFunctionStyleTrailingText: text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_attribute_form_literal_close_tag_is_preserved(self): + # Attribute form ends at the LAST , so a literal close inside code survives. + text = ( + '' + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): + # A closed zero-param call is valid; strict mode must not treat it as truncated. + assert _only('') == {"name": "ping", "arguments": {}} + # A no-arg call that never closes is still rejected as truncated. + assert parse_tool_calls_from_text('', allow_incomplete = False) == [] + class TestParityWithJsonStyle: def test_json_tool_call_with_trailing_prose_is_accepted(self): @@ -176,6 +192,37 @@ class TestGemmaNativeStyle: } +class TestLlama3PythonTagStrict: + def test_closed_dot_call_is_accepted(self): + text = '<|python_tag|>get_weather.call(location="Tokyo")' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert json.loads(calls[0]["function"]["arguments"]) == {"location": "Tokyo"} + + def test_truncated_dot_call_is_rejected(self): + # No closing paren (depth > 0 at EOF): truncated, reject in strict mode. + text = '<|python_tag|>get_weather.call(location="Tokyo"' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers it. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestMistralArrayStrict: + def test_closed_array_is_accepted(self): + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}]' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + def test_unclosed_array_is_rejected(self): + # Missing the closing ]; strict mode must not heal it. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers the object by hand. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): text = "cats" @@ -197,3 +244,822 @@ class TestHealingPathUnaffected: assert text[span[0] : span[1]] == ( "cats" ) + + def test_wrapperless_fallback_calls_carry_spans(self): + # The wrapperless fallback must report spans so consumers strip exactly the markup. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + closed = "before cats after" + calls, spans = parse_with_spans(closed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert closed[span[0] : span[1]] == ( + "cats" + ) + + healed = "x dogs" + calls, spans = parse_with_spans(healed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "dogs"} + (span,) = spans + assert healed[span[0] : span[1]] == "dogs" + + +class TestParserLinearity: + """Llama-3 ``.call`` kwargs and Mistral-array healing must stay linear (a regex-per-offset blew up on long truncated bodies).""" + + def test_llama3_unterminated_call_arg_is_linear(self): + import time + + text = '<|python_tag|>upload.call(data="' + "A" * 200_000 # no closing quote/paren + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_huge_wordrun_call_arg_is_linear(self): + import time + + text = "<|python_tag|>upload.call(" + "a" * 200_000 # giant word run, no '=' + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_mistral_unclosed_array_open_braces_is_linear(self): + import time + + text = "[TOOL_CALLS] [" + "{" * 200_000 # unclosed array, all open braces + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_call_kwargs_still_parse(self): + text = '<|python_tag|>do.call(s="hi 😀", n=42, f=1.5, b=true, z=null)' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "s": "hi 😀", + "n": 42, + "f": 1.5, + "b": True, + "z": None, + } + + def test_llama3_call_scientific_notation_args_parse(self): + # Scientific notation must decode as float (the old regex truncated 1e-3 -> 1). + text = "<|python_tag|>calc.call(x=1e-3, y=-2E+4, z=0.5e2, n=42)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"x": 1e-3, "y": -2e4, "z": 50.0, "n": 42} + assert isinstance(args["n"], int) and isinstance(args["x"], float) + + def test_mistral_unclosed_array_recovers_top_level_objects(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"k":1}},' + '{"name":"b","arguments":{"j":2}}' # missing closing ] + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + + +class TestLlamaBuiltinChainAndNesting: + """Llama-3 ``.call`` built-ins: ``; `` chaining and nested-tag isolation.""" + + def test_semicolon_chained_builtin_calls_all_parse(self): + # Only the first call is anchored to <|python_tag|>; the rest chain via ';'. + text = "<|python_tag|>alpha.call(x=1); beta.call(y=2); gamma.call(z=3)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["alpha", "beta", "gamma"] + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_nested_python_tag_in_json_string_arg_is_not_a_call(self): + # A <|python_tag|> literal inside a code arg is data: the outer "python" call wins. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"<|python_tag|>os.call(\'rm -rf /\')"}}' + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "python" + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "<|python_tag|>os.call('rm -rf /')" + + def test_single_builtin_call_unchanged(self): + text = '<|python_tag|>web_search.call(query="cats")' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + +def test_strip_leading_bare_json_call_drops_complete_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A complete Llama-3.2 bare-JSON call is removed; trailing prose is kept. + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"cats"}}') == "" + assert ( + strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') == "done" + ) + + +def test_strip_leading_bare_json_call_drops_truncated_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A truncated call (no closing brace) collapses to "" -- nothing recoverable. + assert ( + strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"weather in S') + == "" + ) + + +def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # No "name" key -> plain JSON answer, left untouched. + assert ( + strip_leading_bare_json_call('{"result": 42, "ok": true}') == '{"result": 42, "ok": true}' + ) + # Prose before the brace -> not a leading bare call, untouched. + assert strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + # Ordinary text untouched. + assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." + + +def test_bare_json_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + alice = '{"name":"Alice","parameters":{"age":30}}' + real = '{"name":"web_search","parameters":{"query":"cats"}}' + # With an enabled set, markerless JSON whose name is not a tool is NOT a call. + assert parse_tool_calls_from_text(alice, enabled_tool_names = {"web_search"}) == [] + # A real call (enabled name) still parses. + got = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in got] == ["web_search"] + # No enabled set (None) keeps the name-agnostic behaviour for direct callers. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == ["Alice"] + # Marker-based forms are NOT gated (an explicit signal is a real call attempt). + xml = '{"name":"Alice","arguments":{}}' + assert parse_tool_calls_from_text(xml, enabled_tool_names = {"web_search"}) + + +def test_strip_leading_bare_json_call_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + alice = '{"name":"Alice","parameters":{"age":30}}' + # Not an enabled tool -> ordinary JSON answer, kept verbatim. + assert strip_leading_bare_json_call(alice, {"web_search"}) == alice + # Enabled tool -> a real call, stripped (trailing prose kept). + assert ( + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"q":1}} hi', {"web_search"} + ) + == "hi" + ) + + +def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): + from core.inference.tool_call_parser import strip_tool_markup + + # Strip uses the LAST so a literal in a value survives; calls strip independently. + text = 'print("") done' + assert strip_tool_markup(text, final = True) == "done" + two = ( + "a 1 mid " + "2 end" + ) + assert strip_tool_markup(two, final = True) == "a mid end" + + +def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): + from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup + + # A literal opener inside a value is data: the strip keeps " done". + text = 'print("") done' + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" + assert strip_tool_markup(text, final = True) == "done" + # Non-final (streaming) keeps an unclosed call buffered, does not eat prose early. + open_text = 'pre print("")' + assert strip_tool_markup(open_text, final = False) == open_text + + +def test_final_strip_removes_magistral_think_reasoning(): + from core.inference.tool_call_parser import strip_tool_markup + + # Magistral reasoning is [THINK]...[/THINK]; end-of-turn must drop it. + text = "[THINK]The user greeted me, I should say hi.[/THINK]Hello! How can I help?" + assert strip_tool_markup(text, final = True) == "Hello! How can I help?" + # A [TOOL_CALLS] living inside the reasoning goes with it. + with_call = '[THINK]Maybe I should search.[/THINK][TOOL_CALLS]search{"q":"x"}' + assert strip_tool_markup(with_call, final = True) == "" + + +def test_streaming_strip_keeps_magistral_think_buffered(): + from core.inference.tool_call_parser import strip_tool_markup + + # Mid-stream (final=False) leaves the reasoning block intact; only end-of-turn removes it. + text = "[THINK]still thinking" + assert strip_tool_markup(text, final = False) == text + + +def test_final_strip_leaves_non_magistral_bracket_text_untouched(): + from core.inference.tool_call_parser import strip_tool_markup + + # Only a LEADING [THINK] block is reasoning; unrelated bracketed prose stays. + text = "See [THINK about it] later" + assert strip_tool_markup(text, final = True) == "See [THINK about it] later" + + +def test_strip_leading_bare_json_call_ignores_nested_name(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A nested "name" must NOT gate the strip; the JSON answer is kept verbatim. + nested_trunc = '{"result":{"name":"web_search","age":' + nested_full = '{"result":{"name":"web_search","age":1}}' + assert strip_leading_bare_json_call(nested_trunc, {"web_search"}) == nested_trunc + assert strip_leading_bare_json_call(nested_full, {"web_search"}) == nested_full + # A real top-level call (even with a top-level array before the name) still strips. + assert ( + strip_leading_bare_json_call( + '{"data":[1,2],"name":"web_search","parameters":{}}', {"web_search"} + ) + == "" + ) + + +def test_mistral_single_object_call_is_stripped_for_display(): + from core.inference.tool_call_parser import ( + _strip_mistral_closed_calls, + parse_tool_calls_from_text, + ) + + # The parser accepts single-object [TOOL_CALLS]{...}, so the strip must remove it too. + text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] + assert _strip_mistral_closed_calls(text) == " tail" + # A literal [TOOL_CALLS] in prose (no following object) is left untouched. + assert _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") == "See the [TOOL_CALLS] docs" + + +def test_tool_call_parser_declares_future_annotations_for_py39_import(): + # PEP 604 X | None annotations need `from __future__ import annotations` on py3.9; guard it stays. + from pathlib import Path + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" + ).read_text() + assert "from __future__ import annotations" in src + + +def test_bare_json_function_alias_parses_and_strips_symmetrically(): + # The "function" alias for the call name must parse and strip symmetrically. + from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_leading_bare_json_call, + _top_level_bare_json_name, + ) + + enabled = {"web_search"} + text = '{"function":"web_search","parameters":{"query":"cats"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = enabled) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert strip_leading_bare_json_call(text, enabled) == "" + + # "name" still takes precedence when both are present; nested aliases are data. + assert _top_level_bare_json_name('{"function":"foo","name":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"function":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"result":{"function":"web_search"}}') is None + # A non-enabled function-alias object is ordinary content and is preserved. + assert ( + strip_leading_bare_json_call('{"function":"not_a_tool","parameters":{}}', enabled) + == '{"function":"not_a_tool","parameters":{}}' + ) + + +class TestMistralOuterOverXmlLiteral: + """Quoted tool XML inside a [TOOL_CALLS] call's arguments is data; the outer call executes. Reverse order keeps the XML.""" + + def test_mistral_v11_arg_quoting_function_xml(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"literal ' + '1"}' + ) + for strict in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = not strict) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_mistral_array_arg_quoting_tool_call_json(self): + text = ( + '[TOOL_CALLS][{"name":"web_search","arguments":{"query":' + '"see {\\"name\\":\\"evil\\"}"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_outer_keeps_winning_over_mistral_literal(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"docs say [TOOL_CALLS]evil[ARGS]{}"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestHealerSignalAlignment: + """The healer buffers only promotable formats; Mistral/Llama text calls stream through.""" + + def test_heal_signals_subset_of_promotable_formats(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", "evil.call(x=1)"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what is <|python_tag|>evil.call(x=1)" + + +class TestPythonTagOuterOverXmlLiteral: + """A leading Llama-3 ``<|python_tag|>`` call owns the turn: tool XML/Mistral + markup quoted in a ``.call(...)`` string argument (or in trailing prose) is + data, so the outer call executes -- parity with the bare-JSON / Mistral / + attribute-form leading-ownership rules. XML before the tag keeps normal order.""" + + def test_call_arg_quoting_complete_function_xml(self): + # A closed in a .call() code arg must not beat the leading python_tag call. + text = ( + '<|python_tag|>python.call(code="' + '1")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "1" + + def test_call_arg_quoting_bare_function_tag_in_query(self): + # A query mentioning must search, not execute a phantom tool. + text = '<|python_tag|>web_search.call(query="how do I use in llama")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "how do I use in llama" + + def test_call_arg_quoting_tool_call_json(self): + text = ( + "<|python_tag|>save_file.call(content=" + '"{\\"name\\": \\"delete\\", \\"arguments\\": {}}")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["save_file"] + + def test_json_form_code_arg_quoting_function_xml(self): + # JSON emission: a in the code arg is data; the outer "python" call runs. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"ls"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "ls" + + def test_call_arg_quoting_mistral_trigger(self): + text = '<|python_tag|>web_search.call(query="see [TOOL_CALLS]evil[ARGS]{}")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_call_wins_over_trailing_xml(self): + # A leading python_tag call owns the turn even when a real XML literal follows. + text = ( + '<|python_tag|>web_search.call(query="cats") ' + "1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_before_python_tag_keeps_xml_order(self): + # A foreign signal BEFORE the tag keeps normal document order (XML wins). + text = ( + "x " + '<|python_tag|>python.call(code="y")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestBareJsonOuterOverXmlLiteral: + """Quoted tool XML inside a leading bare-JSON call is data; XML before the JSON keeps normal order.""" + + def test_bare_json_code_arg_quoting_function_xml(self): + text = ( + '{"name": "python", "arguments": ' + '{"code": "run() # ls"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "run() # ls" + + def test_bare_json_outer_unrestricted_mode(self): + text = '{"name": "python", "parameters": {"code": "ls"}}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + + def test_xml_before_json_keeps_xml_order(self): + text = ( + "cats" + ' {"name": "python", "arguments": {"code": "x"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMagistralThinkRehearsal: + """A call rehearsed inside [THINK]...[/THINK] is reasoning; the real call after wins, and parse agrees with strip.""" + + def test_function_xml_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]I could emit {"query":"x"}' + ' here[/THINK][TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_hermes_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]maybe {"name":"web_search","arguments":' + '{"query":"x"}}[/THINK]' + '[TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_unclosed_think_parses_nothing(self): + text = '[THINK]let me try {"query":"x"}' + assert parse_tool_calls_from_text(text) == [] + + +class TestDisabledBareJsonLiteralNotPromoted: + """A leading non-enabled-name object is content: nothing inside promotes, and a call after it still parses.""" + + def test_literal_inside_disabled_json_stays_data(self): + text = ( + '{"name": "Alice", "note": "try ' + 'x"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_python_tag_literal_inside_disabled_json_stays_data(self): + text = '{"name": "Alice", "note": "<|python_tag|>web_search.call(query=1)"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_disabled_json_still_parses(self): + text = ( + '{"name": "Alice", "note": "x"} ' + '{"name": "web_search", "arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMistralLiteralInsideLeadingJson: + """A [TOOL_CALLS] literal quoted inside a leading JSON object must not be promoted over it.""" + + def test_outer_json_call_wins_over_mistral_literal(self): + text = '{"name": "python", "arguments": {"code": "[TOOL_CALLS]web_search{}"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "[TOOL_CALLS]web_search{}" + + def test_disabled_outer_json_keeps_mistral_literal_as_data(self): + text = '{"name": "Alice", "note": "[TOOL_CALLS]web_search{}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaWrappedWhitespace: + """Whitespace drift around ``call``/``:`` in wrapped Gemma calls must still parse (no fallback exists).""" + + def test_space_after_call_colon_parses(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_space_around_colon_parses(self): + text = '<|tool_call>call : web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_strict_mode_still_requires_the_closing_tag(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestGemmaDottedArgumentKeys: + """Dotted Gemma keys (namespaced schemas) must survive key-quoting or the call is lost.""" + + def test_dotted_key_parses(self): + text = '<|tool_call>call:web_search{user.name:<|"|>bob<|"|>, query:<|"|>x<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"user.name": "bob", "query": "x"} + + +class TestLeadingMistralCallOwnsTheTurn: + """A leading Mistral call wins in document order over literal XML in trailing prose.""" + + def test_leading_mistral_wins_over_trailing_xml_literal(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"cats"} ' + "Note: 1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_leading_keeps_normal_order(self): + text = ( + "x " + "[TOOL_CALLS]evil[ARGS]{}" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGemmaDottedKeyAfterBareValue: + def test_dotted_key_after_bare_value_is_a_boundary(self): + text = "<|tool_call>call:web_search{query:foo,user.name:bob}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "foo", "user.name": "bob"} + + +class TestNamelessLeadingJsonAnswerIsData: + """A nameless leading JSON answer is an envelope: quoted markup stays data, and a call after it parses.""" + + def test_xml_literal_inside_json_answer_stays_data(self): + text = '{"answer": "use x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_json_answer_still_parses(self): + text = ( + '{"answer": "docs"} {"name": "web_search", ' + '"arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingBareJsonOwnsTurnOverTrailingXml: + """Document order: a leading closed bare-JSON call owns the turn even when + tool XML appears AFTER it (inside-or-after, mirroring the Mistral rule).""" + + def test_leading_call_wins_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}} Example: ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "first"} + + def test_chained_leading_calls_win_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"q":"second"}} ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + def test_non_call_leading_object_defers_to_trailing_real_call(self): + # Nameless/disabled-name objects decline: dropped, and the real trailing call still parses. + for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): + text = lead + ' {"name":"delete_all","arguments":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], (lead, calls) + + def test_leading_xml_call_still_wins_over_trailing_bare_json(self): + text = ( + '{"name":"delete_all","arguments":{}} ' + 'Example: {"name":"lookup","parameters":{"q":"x"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +class TestProseCloseTagAfterClosedFunctionCall: + """A literal in prose after a closed call is data: the call + ends at its first close that is not parameter data, so arguments never + swallow the prose between the real close and the literal.""" + + def test_arguments_do_not_swallow_prose(self): + text = ( + "cats" + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_literal_close_inside_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_arguments_do_not_swallow_prose(self): + # The attribute form shares the first-balanced-close rule: prose closes never fold in. + text = ( + 'cats' + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_attribute_form_literal_close_in_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_two_calls_both_parse(self): + text = ( + 'cats' + 'x=1' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "python"}) + assert [c["function"]["name"] for c in calls] == ["web_search", "python"], calls + + +class TestEnabledNameJsonAnswerIsContent: + """A JSON answer whose top-level name matches an enabled tool but has no + call shape is content: the parser rejects it, so the strip and the drain + gate must keep it visible too.""" + + def test_answer_survives_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + ans = '{"name":"web_search","result":"no call"}' + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_answer_does_not_route_to_draining(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + assert not _looks_like_enabled_bare_json( + '{"name":"web_search","result":"no call"}', {"web_search"} + ) + + def test_real_call_still_strips_and_drains(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + from core.inference.tool_call_parser import strip_leading_bare_json_call + + real = '{"name":"web_search","parameters":{"q":"x"}}' + assert strip_leading_bare_json_call(real, {"web_search"}) == "" + assert _looks_like_enabled_bare_json(real, {"web_search"}) + + def test_arguments_string_call_still_strips(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + call = '{"name":"web_search","arguments":"{\\"q\\":\\"x\\"}"} tail' + assert strip_leading_bare_json_call(call, {"web_search"}) == "tail" + + +class TestAttributeFormLeadingContainment: + """A leading attribute-form call owns the turn: markup quoted inside its + parameter is data, not a call for the shared XML parser to promote.""" + + def test_quoted_tool_call_inside_param_stays_data(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'find ' + '{"name":"delete","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "delete" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_real_xml_call_before_attribute_form_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '{"name":"delete","arguments":{}} Example: ' + 'x' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert calls[0]["function"]["name"] == "delete" + + +class TestParameterKeepsMultipleLiteralCloses: + """A parameter that provably closes with its own tag keeps every literal + function close inside it as data (regression: the first literal close was + treated as ending the parameter, truncating the value).""" + + def test_two_literal_closes_in_one_parameter(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '' + "a b c " + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "a b c" + } + + def test_strip_removes_the_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + text = ( + '' + "a b c after" + ) + assert strip_tool_markup(text, final = True) == "after" + + def test_unclosed_parameter_still_heals_at_function_close(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + calls = parse_tool_calls_from_text( + "val", + enabled_tool_names = {"web_search"}, + ) + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "val"} + + +class TestMistralPreambleOwnership: + """A visible preface before the first Mistral call must not hand the turn + to a later XML literal: the Mistral call is first in document order.""" + + def test_v11_named_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS]web_search[ARGS]{"query":"cats"} Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_array_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS][{"name":"web_search","arguments":{"query":"cats"}}] Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_call_before_trigger_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "1 then " + '[TOOL_CALLS][{"name":"web_search","arguments":{}}]' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert calls[0]["function"]["name"] == "evil" + + def test_prose_mention_without_call_shape_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "See [TOOL_CALLS] docs for details. " + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"evil"}) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestBareJsonStripRequiresTopLevelName: + """The strip's shape gate requires the parser's TOP-LEVEL name in every + mode: a JSON answer with only a nested name is content, even name-agnostic.""" + + def test_nested_name_answer_survives_name_agnostic_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + ans = '{"parameters":{},"result":{"name":"web_search"}}' + assert strip_leading_bare_json_call(ans) == ans + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_real_call_still_strips_name_agnostic(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index c2dc1fe8db..7fe52a664d 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,15 +24,34 @@ import re as _re _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) assert _m, "could not extract _TOOL_XML_RE source" -_ns = {"_re": _re} +# Provide both helpers so the extracted _strip_tool_xml_for_display resolves. +from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls + +_ns = { + "_re": _re, + "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_function_xml_calls": _strip_function_xml_calls, +} exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] + +_xml_helper = _re.search( + r"def _strip_tool_xml\(text: str\) -> str:\n(?: .+\n)+", + _src, +) +assert _xml_helper, "could not extract _strip_tool_xml source" +assert "_strip_mistral_closed_calls" in _xml_helper.group( + 0 +), "extracted _strip_tool_xml no longer runs the Mistral balanced strip" +exec(_xml_helper.group(0), _ns) + _helper = _re.search( r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" r"(?: .+\n)+", _src, ) assert _helper, "could not extract _strip_tool_xml_for_display source" +assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] @@ -46,6 +65,15 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): + # [TOOL_CALLS] with nested JSON needs the Mistral balanced-brace strip, not the regex. + text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "web_search" not in out, out + assert out == "ok tail" + + def test_strips_well_formed_tool_call(): text = ( "Let me search.\n" @@ -73,6 +101,25 @@ def test_strips_function_only_well_formed(): assert "Done." in cleaned +def test_strips_function_attribute_form(): + # Attribute form must strip from the route too; dotted/hyphenated names included. + text = ( + 'Sure.\n\n' + "\nSydney\n\n\nDone." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "Sure." in cleaned and "Done." in cleaned + + dotted = 'A x B' + assert _TOOL_XML_RE.sub("", dotted) == "A B" + + # Auto-Heal-disabled display contract still preserves literal markup. + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "" not in cleaned + + +# Llama-3 <|python_tag|> arm bounds on REAL sentinels only +def test_python_tag_strip_consumes_literal_sentinel_in_arg(): + # A literal <|...|> token inside the arg must not end the strip early. + text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}" + + +@pytest.mark.parametrize( + "sentinel", + [ + "<|eot_id|>", + "<|eom_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + ], +) +def test_python_tag_strip_stops_at_real_sentinel(sentinel): + # A real control sentinel bounds the strip so following text survives. + text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer' + cleaned = _TOOL_XML_RE.sub("", text) + assert ( + cleaned == f"{sentinel}visible answer" + ), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}" + + +def test_python_tag_strip_restarts_on_second_python_tag(): + # A second <|python_tag|> opens a new region; both are stripped. + text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"second python_tag region leaked: {cleaned!r}" + + +def test_route_strip_removes_param_alias_close_tag(): + # Orphan (attribute-form alias of ) must strip too. + assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + assert ( + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + ) + + +def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): + # A literal in a value must not truncate the strip. + text = " tail" + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + + +def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): + # The call ends at its first non-data close; prose after (even a literal ) survives. + from core.inference.tool_call_parser import strip_tool_markup + text = ( + "cats" + " Done. The tag closes a call." + ) + assert strip_tool_markup(text, final = True) == "Done. The tag closes a call." + + +def test_final_strip_keeps_prose_mentioning_bare_markers(): + # A false-alarm marker in prose must not drop trailing text; only call-start-shaped text drops. + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + "See [TOOL_CALLS] docs for details. More prose after.", + "<|python_tag|> is the Llama marker. Explanation continues.", + "The <|tool_call> opener wraps Gemma calls.", + ): + assert strip_tool_markup(text, final = True) == text + # A bare marker at end-of-text is a fragment and still drops. + assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text" + + +def test_final_strip_still_drops_truncated_marker_calls(): + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + '[TOOL_CALLS][{"name":"web_search","argu', + '[TOOL_CALLS]web_search[ARGS]{"q":"x', + '<|python_tag|>{"name":"web_search","par', + '<|python_tag|>foo.call(items=["a', + "<|tool_call>call:web_search{query:tru", + ): + assert strip_tool_markup(text, final = True) == "" + + +def test_chained_bare_json_strip_consumes_all_calls(): + # Next-turn history must not keep an executed call, else it replays. + from core.inference.tool_call_parser import strip_leading_bare_json_call + + enabled = {"web_search", "python"} + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == "" + assert ( + strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled) + == "trailing prose" + ) + # The chain stops at a non-call answer object, which stays visible. + call_then_answer = ( + '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' + ) + assert ( + strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled) + == '{"name":"web_search","result":"data"}' + )