diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7978a200c0..ae4b386589 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml index 9c28e21672..aaf258d615 100644 --- a/.github/workflows/lockfile-audit.yml +++ b/.github/workflows/lockfile-audit.yml @@ -60,11 +60,11 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: '3.12' diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 8186c07211..0bc216d65a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1610,6 +1610,13 @@ jobs: - name: Install Pester v5 shell: pwsh run: | + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } Set-PSRepository PSGallery -InstallationPolicy Trusted Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser Import-Module Pester -MinimumVersion 5.5.0 diff --git a/.gitignore b/.gitignore index 5017a114b8..9337a1d9ee 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ outputs/ exports/ /datasets/ studio/backend/assets/datasets/ +# Generated async worker / reviewer transcripts (never part of the product). +studio/backend/async_task_outputs/ unsloth_training_checkpoints/ *.gguf *.safetensors diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 3b33a7ce84..9806e90666 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -66,6 +66,14 @@ "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, { "package": "cffi", "file": "cffi/setuptools_ext.py", @@ -159,16 +167,16 @@ "file": "fastapi/routing.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", - "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" }, { "package": "fastapi", "file": "fastapi/routing.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", + "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" }, { "package": "fastmcp-slim", @@ -186,6 +194,14 @@ "evidence": "FS: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" + }, { "package": "fastmcp-slim", "file": "fastmcp/server/auth/providers/jwt.py", @@ -255,16 +271,16 @@ "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", - "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" + "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", + "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", - "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" + "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", + "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" }, { "package": "huggingface-hub", diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b4c908b0cb..22a21a2ebc 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,6 +564,12 @@ 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, not a runtime + # binding: the name (`annotations`, ...) is never loaded, so it can never + # "resolve" to a use. Skip it so a legitimately-added future import + # (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged. + 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 +594,23 @@ 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 and must not block: when a name + # keeps its spelling but its import source is moved A -> B in THIS diff (the + # old `from A import x` is removed at module level and a new `from B import x` + # is added), the swap is intentional, not a silent re-point to a pre-existing + # different object. This mirrors the relocation tolerance already applied to + # TARGET-MISSING. The dangerous case -- the name now resolving to a target + # that already existed before (shadow/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/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb79..ad78157418 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation. The default get_inference_backend() returns an InferenceOrchestrator that delegates to a subprocess. The original InferenceBackend runs inside the subprocess and can be imported directly from .inference when needed. + +Public names are resolved lazily (PEP 562): importing this package -- or a +dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull +the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML +backend and its Studio dependencies). Those load only when a public name is +actually accessed, so standalone helpers stay unit-testable without the full +inference stack. """ -from .orchestrator import InferenceOrchestrator, get_inference_backend -from .llama_cpp import LlamaCppBackend - -# Expose InferenceOrchestrator as InferenceBackend for backward compat. -InferenceBackend = InferenceOrchestrator +from typing import TYPE_CHECKING __all__ = [ "InferenceBackend", @@ -21,3 +24,33 @@ __all__ = [ "get_inference_backend", "LlamaCppBackend", ] + +# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator. +_LAZY_ATTRS = { + "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"), + "InferenceBackend": ("orchestrator", "InferenceOrchestrator"), + "get_inference_backend": ("orchestrator", "get_inference_backend"), + "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"), +} + + +def __getattr__(name): + try: + submodule, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + from importlib import import_module + + value = getattr(import_module(f"{__name__}.{submodule}"), attr) + globals()[name] = value # cache so later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) + + +if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names + from .llama_cpp import LlamaCppBackend + from .orchestrator import InferenceOrchestrator, get_inference_backend + InferenceBackend = InferenceOrchestrator diff --git a/studio/backend/core/inference/chat_eos.py b/studio/backend/core/inference/chat_eos.py new file mode 100644 index 0000000000..2a5d0db228 --- /dev/null +++ b/studio/backend/core/inference/chat_eos.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve a chat model's assistant-turn-end stop tokens. + +Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships +config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its +small chat variants ship no generation_config), so generation runs past the turn +and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns. + +Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it +actually uses to end a turn), not raw vocab membership: a base/coder model can +carry ChatML control tokens in a shared vocab without using them, and a loader +may have synced ``eos_token`` to the document terminator. Dependency-light (no +torch / unsloth) so it is unit-testable without the full inference stack. +""" + +from typing import Optional + +# Canonical assistant-turn-end markers per chat family. +_CHAT_TURN_END_TOKENS = ( + "<|im_end|>", # ChatML: Qwen, Yi + "<|eot_id|>", # Llama 3.x + "<|eom_id|>", # Llama 3.x tool turns + "", # Gemma + "", # Gemma-4 + "<|end|>", # Phi + "<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's) +) +# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has +# its own streamer, so its eos is left untouched. +_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>") + + +def _eos_id_set(eos_token_id) -> set: + if isinstance(eos_token_id, (list, tuple)): + return {int(t) for t in eos_token_id if t is not None} + if eos_token_id is not None: + return {int(eos_token_id)} + return set() + + +def _collect_template_text(chat_template) -> str: + """Flatten a tokenizer ``chat_template`` into one scannable string. + + Usually the template is a single jinja string, but multi-variant models + (e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a + ``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list + of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case + would skip turn-end detection for those valid models, so gather every string + leaf (variant names are harmless: they never contain the markers). + """ + if isinstance(chat_template, str): + return chat_template + if isinstance(chat_template, dict): + values = chat_template.values() + elif isinstance(chat_template, (list, tuple)): + values = chat_template + else: + return "" + parts = [_collect_template_text(v) for v in values] + return "\n".join(p for p in parts if p) + + +def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list: + """eos of ``id_tokenizer`` plus any canonical turn-end marker the + ``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` -- + the tokenizer generation actually uses. + + Pass the same tokenizer for both at load time. After a mapped ``get_chat_template`` + pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective + template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped + template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab + folds the turn-end token onto the doc-eos id, and generate_stream re-reads the + original tokenizer, so resolving ids on the mapped tokenizer would store the wrong + (doc-eos) id and let generation run past the real turn marker.""" + ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None)) + template = _collect_template_text(getattr(template_tokenizer, "chat_template", None)) + if not template or any(h in template for h in _HARMONY_MARKERS): + return sorted(ids) + unk = getattr(id_tokenizer, "unk_token_id", None) + for marker in _CHAT_TURN_END_TOKENS: + if marker in template: + try: + tid = id_tokenizer.convert_tokens_to_ids(marker) + except Exception: + tid = None + if tid is not None and tid != unk and int(tid) >= 0: + ids.add(int(tid)) + return sorted(ids) + + +def resolve_chat_turn_end_eos_ids(tokenizer) -> list: + """tokenizer.eos plus any canonical turn-end marker the model's chat_template + actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended + to be resolved once at load. Returns eos unchanged for harmony templates.""" + return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer) + + +def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]: + """Merged eos_token_id list, or None if ``current_eos`` already covers every + resolved turn-end id. Used to repair a model's generation_config at load so + every ``.generate()`` path (vision, tool loops) stops at the turn boundary.""" + if not turn_end_ids: + return None + current_set = _eos_id_set(current_eos) + if set(turn_end_ids) <= current_set: + return None + return sorted(current_set | set(turn_end_ids)) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index b85e9c348a..dfd4c1c0bc 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -3,12 +3,60 @@ """ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg -fallback for templates that reject reasoning/tools args. +fallback for templates that reject reasoning/tools args, plus the shared +native-chat-template fallback used by the transformers and MLX backends. """ +import copy +import json +import logging from typing import Optional +logger = logging.getLogger(__name__) + + +def _normalize_tool_call_arguments(messages: list) -> list: + """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON + string to a dict. + + The OpenAI wire format carries ``arguments`` as a JSON string, but some chat + templates (e.g. the stricter Qwen tool templates shipped with mlx-community + checkpoints) iterate ``arguments.items()`` and raise + ``TypeError: Can only get item pairs from a mapping.`` on the string form + when a prior tool call is re-rendered on the next turn. A dict works on both + strict and lenient templates, so parse the string; leave non-JSON or non-dict + values untouched. Returns the original list unchanged when nothing needed + coercing (no copy).""" + mutated = False + out: list = [] + for msg in messages: + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not tool_calls: + out.append(msg) + continue + new_calls = [] + msg_changed = False + for call in tool_calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except (ValueError, TypeError): + parsed = None + if isinstance(parsed, dict): + call = {**call, "function": {**fn, "arguments": parsed}} + msg_changed = True + new_calls.append(call) + if msg_changed: + out.append({**msg, "tool_calls": new_calls}) + mutated = True + else: + out.append(msg) + return out if mutated else messages + + def apply_chat_template_for_generation( tokenizer, messages: list, @@ -38,21 +86,209 @@ def apply_chat_template_for_generation( attempts.append(dict(reasoning_kwargs)) attempts.append({}) - last_exc: Optional[Exception] = None - for kwargs in attempts: + def _render(msgs: list) -> str: + last_exc: Optional[Exception] = None + for kwargs in attempts: + try: + return tokenizer.apply_chat_template( + msgs, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ) + except TypeError as e: + last_exc = e + continue + except Exception as e: + last_exc = e + break + if last_exc is not None: + raise last_exc + raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + + try: + return _render(messages) + except Exception: + # Strict tool templates reject the JSON-string ``arguments`` form via + # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. + # Original messages render first, so working templates stay byte-identical. + normalized = _normalize_tool_call_arguments(messages) + if normalized is messages: + raise + return _render(normalized) + + +def render_native_template( + *, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: list, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Render ``messages`` + ``tools`` with the model's NATIVE chat template. + + Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit + the ``tools`` schema, so a tool-calling turn silently stops advertising tools. + The native template ships in the model repo and carries the family's + tool-calling syntax. It is loaded straight from the repo (bypassing any + override on the live tokenizer) and cached on ``model_info``. Returns the + rendered prompt only if the native template actually emits the tools (render + differs with vs without tools); otherwise ``None``. + + ``hf_token`` is the token the model was loaded with -- passed to the repo load + so a gated/private model's native template can still be fetched (otherwise the + fallback fails silently and keeps the override prompt that dropped tools). + + ``trust_remote_code`` is sourced from ``model_info`` (the value the model was + actually loaded with) rather than a call-site argument, so the native-template + reload uses exactly the consent already granted at load. A custom-code tokenizer + repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is + passed, so without this the fallback fails silently and keeps the tool-dropping + prompt for a model the user already consented to run remote code for. For a LoRA + adapter the reload targets the base model, whose remote code was gated and loaded + under the same stored flag, so re-passing it executes no unconsented code. + """ + # ``apply_fn`` lets a backend inject its own render; defaults to the module helper. + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + native_tpl = model_info.get("native_chat_template") + if native_tpl is None: + # A LoRA adapter's native template lives on the base model, not the adapter id. + template_source = model_info.get("base_model") or active_model_name + # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can + # instantiate its class (the stored flag already covers template_source). + trust_remote_code = bool(model_info.get("trust_remote_code", False)) try: - return tokenizer.apply_chat_template( - messages, - tokenize = False, - add_generation_prompt = True, - **kwargs, + from transformers import AutoTokenizer + nt = AutoTokenizer.from_pretrained( + template_source, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) - except TypeError as e: - last_exc = e - continue - except Exception as e: - last_exc = e - break - if last_exc is not None: - raise last_exc - raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + native_tpl = nt.chat_template or False + except Exception as exc: + logger.warning( + "Could not load native chat template for '%s': %s", + template_source, + exc, + ) + # A failed fetch is not "no template": leave the sentinel unset so the next + # call retries (caching False would pin the tool-dropping override). + return None + model_info["native_chat_template"] = native_tpl + if not native_tpl: + return None + + tokenizer = model_info.get("tokenizer") or model_info.get("processor") + if tokenizer is None: + return None + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) + # Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the + # generation lock) races concurrent requests. + try: + render_tokenizer = copy.copy(tokenizer) + render_tokenizer.chat_template = native_tpl + except Exception as exc: + logger.warning( + "Could not clone tokenizer for native-template render of '%s': %s", + active_model_name, + exc, + ) + return None + try: + with_tools = apply_fn( + render_tokenizer, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + no_tools = apply_fn( + render_tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "Native-template tool render failed for '%s': %s", + active_model_name, + exc, + ) + return None + return with_tools if with_tools != no_tools else None + + +def render_with_native_template_fallback( + *, + formatted_prompt: str, + tokenizer, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: Optional[list], + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> str: + """Return ``formatted_prompt``, swapping in a native-template render when an + override template dropped the ``tools`` schema. + + If ``tools`` were requested but the live render is identical with and without + them (detected by comparison, robust against tool names in the system prompt), + re-render with the model's native template. Shared by the transformers and MLX + backends so both advertise tools consistently. ``hf_token`` is forwarded so a + gated/private model's native template can still be fetched.""" + if not tools: + return formatted_prompt + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + # Probe whether the live template dropped the schema. A tools-requiring template + # can raise here; on any error keep the valid tools prompt rather than lose it. + try: + probe_no_tools = apply_fn( + tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "No-tools probe failed for '%s'; keeping the existing tools prompt: %s", + active_model_name, + exc, + ) + return formatted_prompt + if formatted_prompt != probe_no_tools: + return formatted_prompt # template already emits the tools schema + native_prompt = render_native_template( + model_info = model_info, + active_model_name = active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = apply_fn, + hf_token = hf_token, + ) + if native_prompt: + logger.info( + "Override template for '%s' dropped tool schemas; using the model's " + "native template for this tool-calling turn.", + active_model_name, + ) + return native_prompt + return formatted_prompt diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4dca4db768..164f202681 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -27,6 +27,10 @@ from utils.hardware import ( from core.inference.audio_codecs import AudioCodecManager from core.inference.runtime_context import runtime_context_length from core.inference.message_content import content_to_text +from core.inference.chat_eos import ( + chat_eos_repair, + resolve_chat_turn_end_eos_ids_using, +) from io import StringIO import structlog from loggers import get_logger @@ -210,6 +214,50 @@ class InferenceBackend: # API uses -1 to disable top-k; transformers uses 0. return 0 if top_k < 0 else top_k + def _resolve_chat_eos(self, model_name: str) -> None: + """Resolve this chat model's assistant-turn-end stop tokens once at load, + cache them in model_info, and repair generation_config so every + ``.generate()`` path stops at the turn boundary. + + Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with + ``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no + ``generation_config.json``, so paths that read ``generation_config`` (the + vision path, tool loops) run past the turn and loop. Turn-end markers are + derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids), + so base/coder models and harmony templates are left untouched. + """ + info = self.models.get(model_name) or {} + model = info.get("model") + container = info.get("tokenizer") + tokenizer = getattr(container, "tokenizer", container) # unwrap processors + if model is None or tokenizer is None: + return + # Vision models carry the chat_template on the processor, not the inner + # tokenizer. Read markers from whichever has one, but resolve ids on the + # generation tokenizer, else the vision path misses the turn-end token. + template_source = container if getattr(container, "chat_template", None) else tokenizer + try: + turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer) + except Exception as e: # never block a load on eos resolution + logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e) + return + info["chat_turn_end_eos_ids"] = turn_end_ids + + gen = getattr(model, "generation_config", None) + if gen is None: + return + repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids) + if repaired is None: + return + previous = gen.eos_token_id + gen.eos_token_id = repaired + logger.info( + "Repaired generation_config.eos_token_id for %s: %s -> %s", + model_name, + previous, + repaired, + ) + def load_model( self, config: ModelConfig, @@ -221,6 +269,9 @@ class InferenceBackend: gpu_ids: Optional[list[int]] = None, ) -> bool: """Load any model: base, LoRA adapter, text, or vision.""" + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. if max_seq_length <= 0: max_seq_length = 2048 @@ -231,6 +282,8 @@ class InferenceBackend: # Already loaded? if model_name in self.models and self.models[model_name].get("model"): logger.info(f"Model {model_name} already loaded") + if hf_token: + self.models[model_name]["hf_token"] = hf_token self.active_model_name = model_name return True @@ -246,6 +299,14 @@ class InferenceBackend: ) self.models[model_name] = { + # Per-model token: the native-template fallback must use the + # token this model was loaded with, not whichever loaded last. + "hf_token": hf_token, + # Per-model consent: the native-template reload must re-use the + # exact trust_remote_code this model (and a LoRA's base) was loaded + # with, so a custom-code tokenizer repo can be re-fetched without + # executing any code the user did not already consent to. + "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, "is_audio": config.is_audio, @@ -496,6 +557,7 @@ class InferenceBackend: max_seq_length, ) + self._resolve_chat_eos(model_name) self._load_chat_template_info(model_name) self.active_model_name = model_name @@ -946,6 +1008,22 @@ class InferenceBackend: tokenizer, chat_template = template_name, ) + # The mapper installs the effective template only now, at generate + # time, so re-resolve and UNION into the load-time cache (never + # overwrite). get_chat_template can return a remapped tokenizer + # (turn-end folded onto doc-eos) while generate_stream reads the + # original, so take marker strings from the mapped template but + # resolve their ids on the original. + try: + _gen_tok = model_info.get("tokenizer") or tokenizer + refreshed = resolve_chat_turn_end_eos_ids_using( + getattr(tokenizer, "tokenizer", tokenizer), + getattr(_gen_tok, "tokenizer", _gen_tok), + ) + existing = model_info.get("chat_turn_end_eos_ids") or [] + model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed)) + except Exception as e: + logger.warning(f"Could not refresh chat turn-end eos after template: {e}") else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" @@ -975,6 +1053,27 @@ class InferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, ) + + # If tools were requested but the (possibly overridden) template ignored + # them, fall back to the model's native template (shared with MLX). + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) + + formatted_prompt = render_with_native_template_fallback( + formatted_prompt = formatted_prompt, + tokenizer = tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = template_messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = self._apply_chat_template_for_generation, + hf_token = model_info.get("hf_token"), + ) + logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") @@ -1382,7 +1481,8 @@ class InferenceBackend: min_p = min_p, repetition_penalty = repetition_penalty, do_sample = temperature > 0, - eos_token_id = tokenizer.eos_token_id, + # Resolved once at load (chat_template-derived turn-end tokens). + eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id, pad_token_id = tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ccfc5cdfe..455d1d084c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -38,9 +38,25 @@ 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 (legacy helper only knew / 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 @@ -8398,6 +8424,10 @@ class LlamaCppBackend: "duration_ms": round((time.monotonic() - started_at) * 1000.0), } + # Enabled-name gate for the markerless Gemma strip (disabled/example + # names stay visible). Set per iteration; None = pre-loop name-agnostic. + _enabled_tool_names = None + def _strip_tool_markup( text: str, *, @@ -8406,11 +8436,22 @@ 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, enabled_tool_names = _enabled_tool_names + ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text + # Shared parser patterns (not the legacy tool_healing set) so textual + # Mistral/python_tag calls entering DRAINING never leak. Balanced strips + # first (nested JSON removed whole); no final trim so length compares hold. + text = _strip_mistral_closed_calls(text) + text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names) + # Parser-accurate scans close at each call's REAL terminator before + # the regex arms: literal markup inside a value is data. + text = _strip_function_xml_calls(text, final = True) + text = _strip_glm_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -8456,6 +8497,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 +8517,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 +8531,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 (not the enlarged range) 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 +8542,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 an ordinary 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 (like safetensors). + _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 +8839,48 @@ class LlamaCppBackend: is_prefix = True break - if is_match: + # Signal-less call shapes (mirror the safetensors + # loop): Llama-3.2 bare {"name":..} and Gemma + # call:NAME{...} would otherwise stream raw. + _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 call: drain + # rather than 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 + elif ( + "call:".startswith(stripped_buf) + or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf) + is not None + or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None + ): + # Whitespace-tolerant like the parser. + if _GEMMA_BARE_TC_RE.match(stripped_buf): + _drain_silently = True + elif len(stripped_buf) < _MAX_BUFFER_CHARS: + _hold_buffer = 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 +8897,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 +8926,17 @@ 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 (the signal-only + # gate below would flush the raw JSON to the user). + _bare_eos = strip_llama3_leading_sentinels(stripped_buf) + # Gate on enabled names so an ordinary 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 +8962,26 @@ class LlamaCppBackend: "text": cumulative_display, } else: + # Held buffer was 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 + # 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. - _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, - ) + # Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less + # calls carry no XML signal, so a signal gate would let them slip past. + _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 +9098,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: we only reach DRAINING when 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 +9112,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 +9130,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 markup to strip and didn't parse. With + # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers 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( @@ -9019,6 +9154,29 @@ class LlamaCppBackend: _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) + # Collapse exact-duplicate calls and cap the count for the TEXTUAL + # fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN). + if tool_calls and not has_structured_tc and len(tool_calls) > 1: + _seen_keys: set = set() + _deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in _seen_keys: + continue + _seen_keys.add(_key) + _deduped.append(_tc) + if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(_deduped) != len(tool_calls): + logger.info( + "GGUF textual fallback: collapsed %d repeated tool call(s) " + "in one turn to %d", + len(tool_calls), + len(_deduped), + ) + tool_calls = _deduped + # disable_parallel_tool_use: execute only the first tool call # this turn. Truncate before building assistant_msg so the # conversation stays consistent and extra calls are never executed. @@ -9144,6 +9302,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 +9327,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 correction 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/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..45f46fef2f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -104,6 +104,9 @@ class MLXInferenceBackend: ) -> bool: import mlx.core as mx + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) @@ -168,11 +171,20 @@ class MLXInferenceBackend: self.active_model_name = model_name self.models[model_name] = { + # Per-model token for the native-template fallback (matches transformers). + "hf_token": hf_token, + # Per-model consent for the native-template reload: re-use the exact + # trust_remote_code this model was loaded with (matches transformers). + "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, "is_vision": is_vision, "is_lora": getattr(config, "is_lora", False), + # For a LoRA adapter the native chat template lives on the base model. + "base_model": getattr(config, "base_model", None) + if getattr(config, "is_lora", False) + else None, "is_audio": False, "audio_type": None, "has_audio_input": False, @@ -355,6 +367,7 @@ class MLXInferenceBackend: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + render_with_native_template_fallback, ) prompt = apply_chat_template_for_generation( @@ -368,6 +381,25 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") + # Same parity fix as the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models keep + # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer, + # so probe and native render share a renderer. (The VLM path renders via the + # processor for image tokens and is intentionally not wired here.) + model_info = self.models.get(self.active_model_name, {}) + prompt = render_with_native_template_fallback( + formatted_prompt = prompt, + tokenizer = self._tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + hf_token = model_info.get("hf_token"), + ) + sampler = make_sampler( temp = temperature, top_p = top_p, diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index c73134b4a2..fe1aca0e4a 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -29,10 +29,26 @@ 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 +# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing) +# actually promotes. The parser module's broader signal list also covers Llama +# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose +# full parser handles them; buffering those here would hold a streamed +# client-tool call until finalization and then flush it as prose (this healer +# cannot promote them), so the passthrough keeps its own 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 +60,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 +214,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 +264,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 +291,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 +524,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..8e86d09754 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -21,12 +21,22 @@ from typing import Callable, Generator, Optional from loggers import get_logger from core.inference.tool_call_parser import ( + _GEMMA_BARE_TC_PREFIX_RE, + _GEMMA_BARE_TC_RE, _TOOL_ALL_PATS, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_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,16 +60,63 @@ 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, not answering; nudge a call. +# Negative lookahead drops negated forms ("I will not") so a refusal doesn't trigger it. 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, not a hardcoded set. Mirrors GGUF tool_hint. +_REPROMPT_INSTRUCTION_TEMPLATE = ( + "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." +) + +# No grammar constraint here (unlike llama-server's lazy grammar): collapse +# exact-duplicate calls and cap the count so a runaway turn cannot fan out. +_MAX_TOOL_CALLS_PER_TURN = 8 + + +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, *, auto_heal_tool_calls: bool = True, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: - """Strip open-ended tool XML from display text without trimming whitespace.""" + """Strip open-ended tool XML from display text without trimming whitespace. + ``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a + disabled/example name in prose is kept (mirrors the parser gate).""" if not (auto_heal_tool_calls or tool_protocol_active): return text + # Mirror the final strip's scan order so streaming and final display agree: + # balanced strips first (nested JSON removed whole), then the guarded + # function-XML/GLM scans that close at each call's REAL terminator, so literal + # markup inside argument values is data and trailing prose survives. No final + # trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK] + # is dropped (bracket form, not the reasoning channel's ); an unclosed + # [THINK] holds until [/THINK] so the cleaned text stays monotonic. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) + text = _strip_function_xml_calls(text, final = True) + text = _strip_glm_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -70,10 +127,11 @@ def _strip_tool_markup_final( *, auto_heal_tool_calls: bool, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: if not (auto_heal_tool_calls or tool_protocol_active): return text - return strip_tool_markup(text, final = True) + return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names) def _status_for_tool(tool_name: str, arguments: dict) -> str: @@ -81,6 +139,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 +264,11 @@ def run_safetensors_tool_loop( kb_search_count = 0 final_attempt_done = False next_call_id = 0 + reprompt_count = 0 + # Real tool-call turns completed. Only turns that actually executed a tool count + # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a + # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. + _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" @@ -215,9 +286,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 +304,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 an ordinary 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 = "" @@ -315,6 +392,7 @@ def run_safetensors_tool_loop( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -345,6 +423,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -367,6 +446,64 @@ def run_safetensors_tool_loop( is_prefix = True break + # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML + # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses + # as a call, else stream as content. Non-call text is always recovered downstream. + 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 instead of leaking the raw prefix; 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. + + # Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry: + # buffer it here or it streams raw until the end-of-turn safety net. + # ``(? len(last_emitted): last_emitted = cleaned @@ -407,6 +545,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -419,44 +558,76 @@ 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 + ): + # A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary + # JSON answer falls through to the else and streams as content, 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, enabled_tool_names = _enabled_tool_names + ) + 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 (the Llama-3.2 bare-JSON form 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 + # in full; route-level 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 @@ -465,6 +636,7 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_tool_names, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", @@ -476,20 +648,25 @@ 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, + enabled_tool_names = _enabled_tool_names, + ) + # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment + # (plain JSON answers are left 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 { @@ -505,10 +682,14 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_tool_names, ) if tool_calls: next_call_id += len(tool_calls) + # Strip a leading bare-JSON call from the kept content 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. @@ -517,6 +698,27 @@ def run_safetensors_tool_loop( yield {"type": "status", "text": ""} return + # Collapse exact-duplicate calls and cap the count (runaway-turn guard). + if tool_calls: + seen_keys: set = set() + deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in seen_keys: + continue + seen_keys.add(_key) + deduped.append(_tc) + if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(deduped) != len(tool_calls): + logger.info( + "Safetensors: collapsed %d repeated tool call(s) in one turn to %d", + len(tool_calls), + len(deduped), + ) + tool_calls = deduped + assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False @@ -634,6 +836,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 +850,11 @@ 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 turns that executed a tool against the cap; a no-op correction turn doesn't + # consume budget so the model gets its nudge and another tool-enabled turn (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..08a6bf418a 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,39 +2,120 @@ # 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) + - ``<｜tool▁calls▁begin｜>...function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`...`` (DeepSeek R1) + - ``<｜tool▁calls▁begin｜>...<｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` (DeepSeek V3 / V3.1) + - ``NAME\\nk\\nv...`` (GLM 4.5 / 4.6 / 4.7) + - ``<|tool_calls_section_begin|>...<|tool_call_begin|>functions.NAME:IDX<|tool_call_argument_begin|>{json}<|tool_call_end|>...`` (Kimi K2) + +Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ +# Lazy annotations keep the standalone python 3.9 import working. +from __future__ import annotations + +import json +import re +from typing import Any, Optional + +# Qwen/Hermes, Qwen3.5 XML and Gemma 4 live in core.tool_healing; this module adds the rest. 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>", + # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<｜tool_calls_begin｜>", + "<｜tool▁calls｜>", + "<｜tool calls begin｜>", + "<｜tool\\_calls\\_begin｜>", + # Kimi K2 / Moonshot. + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", +) -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, - ) +# DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped. +_DEEPSEEK_OPEN_ALT = ( + r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" +) +_DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>" + +# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at +# end-of-turn. ``[\w-]+`` on ```` tracks OpenAI's +# ``^[a-zA-Z0-9_-]{1,64}$`` so hyphenated MCP names parse like built-ins. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + # Span to the real ```` so a literal one inside a value can't truncate the strip. + re.compile( + r'' + r'(?:(?!).)*' + r"", + re.DOTALL, + ), + re.compile(r"<\|tool_call>.*?", re.DOTALL), + re.compile(r"\[TOOL_CALLS\]\s*\[.*?\](?:\s*)?", re.DOTALL), + # Mistral v11+ ``[TOOL_CALLS]name{json}`` (may chain), close at ``}``. + re.compile(r"\[TOOL_CALLS\]\s*[\w\.\-]+\s*(?:\[ARGS\])?\s*\{.*?\}", re.DOTALL), + # DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end. + re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL), + # Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``. + re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL), + # Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS. + re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", 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 a call-shaped start + # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF 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, + ), + # DeepSeek envelopes truncated mid-stream (any opener); same call-shaped lookahead as above. + re.compile( + _DEEPSEEK_OPEN_RE_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*$", + re.DOTALL, + ), + re.compile(r"<｜tool▁call▁begin｜>(?=\s*function|\s*$).*$", re.DOTALL), + # Kimi K2 envelope truncated. + re.compile( + r"<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*$", + re.DOTALL, + ), + # Gemma wrapper-less ``call:NAME{...}`` is handled by ``_strip_gemma_wrapperless_calls`` (enabled-name gate). +] -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 ```` and the attribute form ```` +# (MiniCPM-5, MiniMax-M2); name class ``[\w.\-]+`` lands in group(1) or group(2). +_TC_FUNC_START_RE = re.compile(r'\s*') +# Body ends at ```` (Hermes) or ```` (Qwen3.5 / MiniCPM-5) +# so it stops at the close even when prose follows (else prose leaked into args). +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +# Horizontal whitespace only (``[^\S\n]*``, not ``\s*``) so the wrapping newline + +# first-line indentation survive; ``_trim_param_value`` trims one newline, preserving +# code indentation (SGLang qwen3_coder). +_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 a fixed offset (char after ``<|python_tag|>``) plus the ``; NAME.call(`` +# chain separator; fixed-offset (not a free scan) ignores ``.call(`` inside JSON args. +_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*\(") +# Llama-3 ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay +# linear on a truncated body; finditer retries every offset of a long run (ReDoS). +_LLAMA3_KEY_RE = re.compile(r"\w+") +_LLAMA3_WS_RE = re.compile(r"\s*") +# ints, decimals (1.5, 1., .5) and sci notation; trailing ``(?![\w.])`` stops a token +# like ``1.2.3`` being truncated to ``1.2`` (which would mis-parse the remainder). +_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 them, each followed by a bare name +# plus ``{json}`` (Magistral) or ``[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); llama.cpp distinguishes the two on ``[CALL_ID]`` (common/chat.cpp). +_MISTRAL_CALL_ID_MARKER = "[CALL_ID]" +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside +# that block is chain-of-thought, not a real call. +_MISTRAL_THINK_OPEN = "[THINK]" +_MISTRAL_THINK_CLOSE = "[/THINK]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# DeepSeek markers (full-width pipe U+FF5C, block U+2581); five outer-open variants like llama.cpp. +_DEEPSEEK_BEGIN_RE = re.compile(_DEEPSEEK_OPEN_RE_SRC) +_DEEPSEEK_END = "<｜tool▁calls▁end｜>" +_DEEPSEEK_CALL_BEGIN = "<｜tool▁call▁begin｜>" +_DEEPSEEK_SEP = "<｜tool▁sep｜>" +_DEEPSEEK_CALL_END = "<｜tool▁call▁end｜>" +# R1 wraps args in a ```json fence with a ``function`` prefix; V3/V3.1 do not. +# Scanned with ``str.find`` -- the regex forms are O(N^2) on truncated bodies. +_DEEPSEEK_R1_FUNC_MARKER = "function" + _DEEPSEEK_SEP +_DEEPSEEK_R1_FENCE = "\n```json\n" +_DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_END)) + +# GLM 4.5-4.7: ``NAME[\n]K...``; the lookahead also allows a +# direct ````/```` (4.7 drops the newline, zero-arg calls close at once). +# Name class ``[\w.\-]+`` keeps prose like ``not a call`` unparsed; +# ``{`` stays with the Qwen JSON parser. +_GLM_TC_OPEN_RE = re.compile(r"\s*([\w.\-]+)\s*(?=\n||)") +_GLM_TC_CLOSE = "" +_GLM_ARG_KEY_OPEN = "" +_GLM_ARG_KEY_CLOSE = "" +_GLM_ARG_VAL_OPEN = "" +_GLM_ARG_VAL_CLOSE = "" +# Strings arrive raw, non-strings via tojson; only unambiguous JSON literals decode +# (bare ``42``/``true``/``null`` stay strings). +_GLM_JSON_NUMERIC_RE = re.compile(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?") + +# Kimi K2 / Moonshot (ASCII pipes). Id ``functions.NAME:IDX`` -- strip ``functions.``/``:N`` for the name. +_KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +_KIMI_SECTION_END = "<|tool_calls_section_end|>" +_KIMI_CALL_BEGIN = "<|tool_call_begin|>" +_KIMI_ARG_BEGIN = "<|tool_call_argument_begin|>" +_KIMI_CALL_END = "<|tool_call_end|>" +_KIMI_ID_RE = re.compile(r"^(?:functions\.)?([\w\.\-]+)(?::(\d+))?$") + +# 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 = "" + +# skip_special_tokens strips the wrapper and ``<|"|>`` markers, so streamed Gemma calls +# arrive as bare ``call:NAME{k:v, ...}``; ``(? 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]...[/THINK]`` so a ``[TOOL_CALLS]`` inside + reasoning is not taken as a real call; an unclosed ``[THINK]`` drops from it on.""" + 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 (array, ``name{json}``, + ``name[ARGS]{json}``) via balanced scanning -- a non-greedy ``\\{.*?\\}`` would + truncate at the first ``}`` and lose nested JSON. Unclosed runs are left for + ``final=True`` cleanup.""" + 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 }`` (no name/array): the parser + # accepts it, so the display strip must remove it too (else it leaks). + 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 too, mirroring the array shape, so a + # ``[TOOL_CALLS]name{json}`` tail doesn't leave ```` as content. + if text.startswith("", cursor): + cursor += len("") + return "".join(out) + + +def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace + scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the + strip like the parser gate: a disabled/example name stays visible; ``None`` + strips every closed call.""" + if _whole_content_is_json_value(text): + return text + n = len(text) + out = [] + # Mirror the parse scan: a leading JSON answer's span is data, kept visible. + cursor = _leading_json_value_end(text) or 0 + if cursor: + out.append(text[:cursor]) + while cursor < n: + m = _GEMMA_BARE_TC_RE.search(text, cursor) + if not m: + out.append(text[cursor:]) + break + disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names + brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{`` + # Same boundary scanner as the parser: strip exactly what it consumed. + end = _gemma_body_brace_end(text, brace) + closed = end is not None + next_index = (end + 1) if closed else len(text) + if not closed: + # Unclosed call: drop an enabled call to EOS; keep a disabled/example name as prose. + out.append(text[cursor:] if disabled else text[cursor : m.start()]) + break + if disabled: + # Disabled/example name is prose: keep it whole. + out.append(text[cursor:next_index]) + else: + out.append(text[cursor : m.start()]) + cursor = next_index # already past the matching ``}`` + 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 _glm_value_close( + text: str, + vs: int, + *, + strict: bool = False, +) -> int: + """Index of the ```` that really ends the GLM value at ``vs``: the + first one whose next non-space token is ````, ```` or + end-of-text AND that sits at balanced quote state (an embedded literal pair + like ``print("")`` lives inside a still-open string). + Quote openers are contextual (single quote only after punctuation, so + apostrophes are prose; double quote also at word start), mirroring the Gemma + scanners. If no candidate balances, the first token-valid one wins -- except + in ``strict`` mode (Auto-Heal off), which refuses the in-quote fallback rather + than execute truncated arguments. Returns -1 if unclosed.""" + n = len(text) + search = vs + first_candidate = -1 + quote = "" + prev = ":" + prev_raw = ":" + qpos = vs # quote-state cursor; advanced incrementally to each candidate + while True: + ve = text.find(_GLM_ARG_VAL_CLOSE, search) + if ve < 0: + return -1 if strict else first_candidate + j = ve + len(_GLM_ARG_VAL_CLOSE) + while j < n and text[j] in " \t\r\n": + j += 1 + if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j): + while qpos < ve: + ch = text[qpos] + if quote: + if ch == "\\" and qpos + 1 < ve: + qpos += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + if not ch.isspace(): + prev = ch + prev_raw = ch + qpos += 1 + if not quote: + return ve + if first_candidate < 0: + first_candidate = ve + search = ve + len(_GLM_ARG_VAL_CLOSE) + + +def _strip_glm_calls(text: str, *, final: bool) -> str: + """Strip GLM 4.x calls by scanning to each call's REAL ```` (the one + after the last consumed ````, mirroring ``_parse_glm_tool_calls``), so + a literal ```` inside a value is data. Qwen ``{json}`` has + no NAME token and is left to the regex arms. ``final`` drops a truncated call to + EOS; otherwise it stays buffered.""" + out: list[str] = [] + cursor = 0 + n = len(text) + while True: + m = _GLM_TC_OPEN_RE.search(text, cursor) + if not m: + break + apos = m.end() + close = -1 + while True: + ks = text.find(_GLM_ARG_KEY_OPEN, apos) + tc = text.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = text.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < n and text[vstart] in " \t\r\n": + vstart += 1 + if not text.startswith(_GLM_ARG_VAL_OPEN, vstart): + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + ve = _glm_value_close(text, vs) + if ve < 0: + break # unclosed -- truncated + apos = ve + len(_GLM_ARG_VAL_CLOSE) + if close >= 0: + out.append(text[cursor : m.start()]) + cursor = close + len(_GLM_TC_CLOSE) + continue + # Truncated GLM call (no real close yet). + if final: + out.append(text[cursor : m.start()]) + cursor = n + # Non-final: leave the unclosed call (and any tail) buffered as-is. + break + out.append(text[cursor:]) + return "".join(out) + + +def strip_tool_markup( + text: str, + *, + final: bool = False, + enabled_tool_names: Optional[set] = None, +) -> str: + """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; + ``final=True`` also drops trailing unclosed runs and trims. ``enabled_tool_names`` + gates the markerless Gemma ``call:NAME{...}`` strip so a disabled/example name in + prose is kept (mirrors the parser gate); ``None`` strips every closed call.""" + if final: + # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket + # form is not the ```` the reasoning channel renders. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + if final: + text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) + # Scan-strip the function-XML form (a literal ```` inside a value is + # data). The regex arms below cover the other formats but no-op on function calls here. + text = _strip_function_xml_calls(text, final = final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + text = _strip_glm_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) + + +# A Qwen/Hermes ````/```` envelope whose arguments carry literal +# DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first +# marker so the pre-pass skips it. +_EMBEDDED_MARKER_RE = re.compile( + _DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN) +) +# Covers ```` and the attribute form. ``<|python_tag|>`` is Llama-3's +# envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi +# example is data; the call-shaped lookahead mirrors the ``_TOOL_ALL_PATS`` python_tag arm +# so a bare prose ``<|python_tag|>`` mention isn't treated as one. +_OUTER_ENVELOPE_OPEN_RE = re.compile( + r'|' + r"|<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\())" +) +# CLOSED outer envelopes, each spanning to its REAL final close so a literal +# ````/```` inside a value is data. Wrapped Gemma counts too. +_OUTER_ENVELOPE_CLOSED_PATS = ( + re.compile(r"(?:(?!).)*", re.DOTALL), + _TOOL_CLOSED_PATS[1], + re.compile(r"<\|tool_call>.*?", re.DOTALL), +) + + +def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool: + first_marker = _EMBEDDED_MARKER_RE.search(content) + if first_marker is None: + return False + # A leading bare-JSON or Mistral [TOOL_CALLS] call is an outer envelope too: + # a DS/Kimi marker in its argument strings is data. + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if content.startswith("{", i): + end = _balanced_brace_end(content, i) + if end is not None and i < first_marker.start(): + name = _top_level_bare_json_name(content[i : end + 1]) + if name is not None and (enabled_tool_names is None or name in enabled_tool_names): + # The closed leading call owns the turn: a marker inside it is argument + # data, one after it a trailing example (same rule as the XML envelopes below). + return True + if name is not None and first_marker.start() <= end: + # A disabled-name leading object is prose (can't own the turn), but a marker + # inside its own strings stays data. A marker AFTER it falls through to the pre-pass. + return True + elif content.startswith(_MISTRAL_TRIGGER, i): + end = _mistral_region_end(content, i) + if end is not None and i < first_marker.start(): + return True + # A closed outer call PRECEDING the first marker owns the turn; the pre-pass must + # not steal a trailing example or argument data. + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + m = _pat.search(content) + if m is not None and m.start() < first_marker.start(): + return True + residue = content + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + residue = _pat.sub("", residue) + marker = _EMBEDDED_MARKER_RE.search(residue) + if marker is None: + return True + # A marker still stands; any opener left in the residue is UNCLOSED. One before the + # marker is a truncated outer call holding the marker as data: skip the pre-pass. + opener = _OUTER_ENVELOPE_OPEN_RE.search(residue) + return opener is not None and opener.start() < marker.start() + + +def _mistral_region_end(text: str, idx: int) -> int | None: + """Exclusive end of the balanced ``[TOOL_CALLS]`` call starting at ``idx``, + or ``None`` when truncated/unrecognised (same shapes as the strip scan: + array, single-object, and named ``name [CALL_ID]? [ARGS]? {json}``).""" + 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: a visible preface must not hand + # the turn to a later XML literal (preamble-tolerant, like the + # wrapperless-Gemma guard). Prose that merely mentions 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 tool signal a non-envelope parser would fire on + (XML forms plus ``<|python_tag|>``, which also runs before the Mistral parser).""" + 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() + # DeepSeek/Kimi markers are foreign to a JSON envelope too: a marker inside a leading + # object routes through the same guard (and, if disabled, the drop-and-parse-the-tail + # recursion, so a real call after the object is still reached). + marker = _EMBEDDED_MARKER_RE.search(content) + if marker is not None and (first is None or marker.start() < first): + first = marker.start() + return first + + +def _xml_signal_inside_leading_bare_json(content: str) -> bool: + """True when the first foreign tool signal is a quoted literal inside a + LEADING bare-JSON call object or JSON answer -- data, not a real call + (sibling of ``_xml_signal_inside_leading_mistral``).""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if i >= n or content[i] not in "{[": + return False + if content[i] == "[": + # A leading array is only ever a structured answer; its literals are data. + end = _balanced_bracket_end(content, i) + if end is None: + return False + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + return first_xml is not None and i < first_xml < end + end = _balanced_brace_end(content, i) + if end is None: + return False + if _top_level_bare_json_name(content[i : end + 1]) is None: + # A NAMELESS object that parses as real JSON is a structured answer / envelope too: + # quoted markup is data, and the decline path drops it and parses the tail. + # 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 (its parser runs first). + 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 data; after the closed object the + # leading call still owns the turn (mirrors the leading-Mistral rule). + return first_xml is not None and i < first_xml + + +def _signal_inside_leading_wrapperless_gemma( + content: str, enabled_tool_names: Optional[set] +) -> bool: + """True when the first foreign tool signal is a quoted literal inside (or + after) a LEADING enabled wrapper-less Gemma call (sibling of the + Mistral/bare-JSON leading guards). Markerless form, so gated on an enabled + name (``None`` keeps the name-agnostic behaviour).""" + first = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a Gemma call too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return False + # A preamble before ``call:NAME{...}`` is normal; what matters is an ENABLED balanced + # call beginning before the first foreign signal. + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return False + if enabled_tool_names is not None and m.group(1) not in enabled_tool_names: + cursor = m.end() + continue + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + return False + if m.end() - 1 < first <= end: + return True + # An enabled call that CLOSES before the signal still owns the turn (inside-or-after + # rule, as for closed bare-JSON/Mistral envelopes), gated on an enabled name. + return enabled_tool_names is not None and end < first + + +def _disabled_gemma_call_end_containing_signal( + content: str, enabled_tool_names: Optional[set] +) -> int | None: + """End offset (exclusive) of the earliest DISABLED wrapper-less Gemma call + whose balanced body contains the first foreign signal, else None. A disabled + name is prose, so the quoted literal is data: the caller drops the span and + recurses on the tail. An ENABLED call defers to the enabled-call guard.""" + if enabled_tool_names is None: + return None + first = _first_foreign_tool_signal(content) + # Mirror the enabled-call guard: the Mistral trigger is foreign here too. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return None + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return None + if m.group(1) in enabled_tool_names: + return None + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + cursor = m.end() + continue + if m.end() - 1 < first <= end: + return end + 1 + cursor = end + 1 + + +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 so calls are never double-counted. + + ``allow_incomplete=True`` (default) heals truncated calls (missing close tag / + unclosed parameter); ``False`` accepts only well-formed closed calls (trailing + prose tolerated), matching llama-server's strict path when Auto-Heal is off. + + ``enabled_tool_names`` gates only the markerless Llama-3.2 bare-JSON form (the + marker-based forms carry an explicit signal, so a disabled-tool name there is a + real call attempt). ``None`` keeps the name-agnostic behaviour.""" + # Drop Magistral [THINK]...[/THINK] BEFORE dispatch: a rehearsed call inside it must + # never be promoted, and the parse path must agree with the display strip. + content = _strip_mistral_reasoning(content) + + # A leading bare-JSON value is decided FIRST: a string argument quoting tool markup + # (XML or a Mistral trigger) must stay data, so the bare-JSON parser takes the outer + # call before any other pass. Precedes the Mistral guard, whose preamble tolerance + # would otherwise 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 content. Drop it and parse the tail. + i = 0 + while i < len(content) and content[i] in " \t\n\r": + i += 1 + # The guard guarantees a balanced leading value (object or array). + end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i) + return parse_tool_calls_from_text( + content[end + 1 :], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A leading enabled wrapper-less Gemma call is decided BEFORE the Mistral guard: its + # body reads as prose to the preamble tolerance below, so a quoted [TOOL_CALLS] would + # otherwise steal the turn. + if _signal_inside_leading_wrapperless_gemma(content, enabled_tool_names): + calls = _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls + + # A DISABLED wrapper-less Gemma call is prose: drop the span and parse the tail BEFORE + # the Mistral guard, whose preamble tolerance would otherwise parse a quoted trigger. + _prose_end = _disabled_gemma_call_end_containing_signal(content, enabled_tool_names) + if _prose_end is not None: + return parse_tool_calls_from_text( + content[_prose_end:], + 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 trailing prose is not promoted over it, nor does a prose preface forfeit 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 + + # DeepSeek/Kimi markers are unique, so try them first -- unless an outer envelope + # opens before the first marker (then the marker is argument data). + if not _marker_inside_leading_envelope(content, enabled_tool_names): + # Dispatch by earliest opener so a quoted DS example inside a Kimi call (or vice + # versa) can't hijack the turn via fixed parser order. + _ds = _DEEPSEEK_BEGIN_RE.search(content) + _ds_pos = _ds.start() if _ds else len(content) + _km_section = content.find(_KIMI_SECTION_BEGIN) + _km_bare = content.find(_KIMI_CALL_BEGIN) + _km_pos = min(p for p in (_km_section, _km_bare, len(content)) if p >= 0) + pre_pass = [ + (_ds_pos, _parse_deepseek_tool_calls), + (_km_pos, _parse_kimi_tool_calls), + ] + pre_pass.sort(key = lambda pair: pair[0]) + for _pos, parser in pre_pass: + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # A leading MiniCPM/MiniMax attribute-form call owns the turn: tool_healing doesn't know + # the wrapper, so a quoted in its parameter would beat + # the outer call. Any earlier signal 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 go through the shared tool_healing + # parser (strict/Auto-Heal contract + nested-marker, trailing-prose, and + # ``<|"|>`` quoted-string handling the GGUF path 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; these run only after it finds + # nothing, so a strict-rejected call is never re-healed here. Blank any + # JSON/Gemma marker coverage first: markup inside a marker's span (even one + # that failed to parse) is that call's data, not a sibling, so a nested + # ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + fallback_content = content + coverage = _tool_healing.marker_coverage(content) + if coverage: + chars = list(content) + for cov_start, cov_end in coverage: + for i in range(cov_start, min(cov_end, len(chars))): + chars[i] = " " + fallback_content = "".join(chars) + for parser in ( + _parse_glm_tool_calls, # GLM 4.x name + _parse_function_xml, # attribute form + _parse_llama3_python_tag, # Llama-3 <|python_tag|> + _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] + ): + calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # Llama-3.2 bare ``{"name":..., "parameters":...}`` (strict shape). Only a LEADING call + # object matches and owns the turn, so an enabled ``call:NAME{...}`` in its arguments + # stays data (Gemma never starts ``{``). + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + + # Gemma wrapper-less ``call:NAME{...}``: markerless, so the same enabled-name gate applies. + return _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + 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 JSON body that never closed its ```` + # is a truncated call, not a finished one. Trailing prose after the close + # is still tolerated (matches the GGUF strict path). + 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 ``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 one wrapping newline the template adds around an XML parameter value + (``\nVALUE\n``), preserving inner indentation. + ``str.strip()`` destroyed code/diff indentation; SGLang's qwen3_coder trims only + the wrapping newline.""" + 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`` sits inside an unclosed ````/```` block -- + i.e. a ```` / ```` opener at ``pos`` is a literal inside an + argument value (e.g. code that prints tool-call XML), not a real nested call. + Compares the last parameter opener before ``pos`` against the last + parameter/function close before it.""" + 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: while it closes after ``pos`` the position is + # argument data, even across several literal function closes. Only an unclosed + # parameter (heal mode) falls back to the first function 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 is promoted to 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 a code/search argument is skipped as data, and + # prose after the real close isn't folded into the last argument (mirrors + # _strip_function_xml_calls and tool_healing._func_close_index). + 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: an unclosed function call 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 + # A ```` opener inside an open parameter 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: a dangling parameter means the call was cut off; a closed + # zero-parameter call stays valid. + 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 (string/number/true/false/null) at ``body[p:]``. + Returns ``(value, consumed_len)`` or ``(None, None)`` if none matches.""" + 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) + # Scientific notation (1e-3, -2E+4, 0.5e2) and decimals decode as float; a bare + # integer stays int. ``"." in v`` alone missed the exponent forms (1e-3 -> 1). + 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]: + """``k=v, ...`` kwargs from a ``.call(...)`` body, left to right (later keys win). + Linear hand-scan replacing the quadratic ``_LLAMA3_KV_RE.finditer`` walk.""" + 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 the Llama-3 emissions: ``<|python_tag|>NAME.call(...)`` (built-in), + ``<|python_tag|>{"name":..., "parameters":...}`` (custom), multi-call via + ``; ``, ``parameters`` or ``arguments`` key.""" + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` and optionally + # ``; ``-chained within one emission. Anchoring to the tag boundary (not a free scan) + # keeps a literal ``<|python_tag|>x.call(...)`` quoted in a custom-form JSON argument + # from being mistaken for a real built-in call. + 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(...)`` with no closing paren: reject in strict mode + # instead of executing a partial. + 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 multiple + # ``; ``-separated objects from one emission. + 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", {}) + # Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value. + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + cursor = brace + end_offset + continue + 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 role label the +# template inserts between ``<|start_header_id|>`` and ``<|end_header_id|>``. +_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 special-token sentinels (and the role label after + ``<|start_header_id|>``) that can leak from a prior turn before a bare-JSON tool + call. Shared by the parser and the streaming buffering guards so a + sentinel-prefixed ``{"name":...}`` is recognised the same everywhere.""" + 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 ``custom_tools`` bare ``{"name":.., "parameters":{..}}`` (no ``<|python_tag|>``), + strict so prose/echoes don't fire. ``enabled_tool_names`` gates on the parsed name so an + ordinary JSON answer isn't misread as a call to a disabled tool; ``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: treat it as a call only when the name is an enabled + # tool, else it is an ordinary JSON answer. + 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 + # JSON-string of one (OpenAI). Looser would fire on ``{"name":"x","parameters":"sentence"}``. + 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 all Mistral emissions: pre-v11 ``[TOOL_CALLS][...]`` / ``[TOOL_CALLS]{...}`` + and v11+ ``[TOOL_CALLS]name{json}`` / ``[TOOL_CALLS]name[ARGS]{json}``.""" + 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 without a ``name`` so v11+ still runs. + 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}`` after each trigger. + 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 a truncated call. In strict mode reject it + # instead of recovering objects by hand below. + 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 + # ``{...}`` instead of re-scanning from every ``{`` (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 (sibling paths and + # SGLang's base detector alias it) so an array object keyed on it keeps args. + 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 _whole_content_is_json_value(text: str) -> bool: + """True when the entire content is one valid JSON value (a structured + answer, e.g. a response_format turn). Markerless scans must treat text + inside it as data: an answer documenting an enabled tool's syntax must + not execute that tool or have the example stripped from display.""" + t = text.strip() + if t[:1] not in "{[": + return False + try: + json.loads(t) + except ValueError: + return False + return True + + +def _leading_json_value_end(text: str) -> int | None: + """End index (exclusive) of a balanced LEADING JSON value that parses as + JSON: a structured answer possibly followed by prose. Markerless scans treat + its contents as data (extends ``_whole_content_is_json_value``); leading-keyed, + so a JSON blob mid-prose is not an answer span.""" + i = 0 + n = len(text) + while i < n and text[i].isspace(): + i += 1 + if i >= n or text[i] not in "{[": + return None + end = (_balanced_brace_end if text[i] == "{" else _balanced_bracket_end)(text, i) + if end is None: + return None + try: + json.loads(text[i : end + 1]) + except ValueError: + return None + return end + 1 + + +def _parse_gemma_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``, plus the + ``skip_special_tokens`` stream where the wrapper and string markers were + stripped (bare ``call:NAME{k:v, ...}``). + + ``enabled_tool_names`` gates on the parsed name: the wrapper-less shape is + indistinguishable from prose documenting the syntax, so a disabled/example + name must not be stolen as a call. ``None`` keeps the name-agnostic behaviour.""" + out: list[dict] = [] + # The WRAPPED form (strict + nested-marker handling) is tool_healing's, which runs + # first: defer content with a wrapped opener. A marker literal alone is not enough -- + # a wrapper-less call mentioning ``<|tool_call>`` would be lost if deferred. + if _GEMMA_TC_RE.search(content): + return out + # A whole-content JSON value is a structured answer: quoted examples must not become calls. + if _whole_content_is_json_value(content): + return out + # Manual cursor: resume AFTER each consumed balanced body so a nested ``call:OTHER{...}`` + # in an argument is never re-matched. A leading JSON answer's span is data -- scan after it. + cursor = _leading_json_value_end(content) or 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None: + break + name = m.group(1) + body_start = m.end() - 1 + end = _gemma_body_brace_end(content, body_start) + if end is None: + # Unclosed call: nothing parseable follows (mirrors the strip contract); + # scanning on would promote quoted argument text. + break + cursor = end + 1 + # Markerless: a disabled/example name is prose, not a call. + if enabled_tool_names is not None and name not in enabled_tool_names: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_stripped_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 + + +def _gemma_body_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` closing the wrapper-less Gemma body at ``brace_pos``. + + Values are raw after ``skip_special_tokens``, so quoted strings (single or + double) hide braces; the quote rules mirror ``_gemma_parse_stripped_body`` so + the boundary always agrees with the body parser. Contextual openers: a single + quote opens only at value-start context (after ``:{[(,=`` -- apostrophes in + ``what's the weather`` are prose), a double quote also at word start (so + ``query:find "a, b"`` hides its delimiters).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + quote = "" + prev = "" + prev_raw = "" + i = brace_pos + n = len(text) + while i < n: + ch = text[i] + if quote: + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + if not ch.isspace(): + prev = ch + prev_raw = ch + 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, name wins) of a bare-JSON object, else None. + + Skips nested objects/arrays so a nested ``"name"`` isn't mistaken for the call name; a + truncated tail returns None so the caller keeps the text.""" + 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 the object with no top-level ``"name"``: fall back to a recorded ``"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"`` aliases the call name. Record it but keep scanning: a top-level + # ``"name"`` still 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 can't prove a top-level name + # exists, so return 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"`` (e.g. {"result":{"name":"web_search",...}}) is data, not the + # call name, so it must not gate the strip. An un-extractable 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 (dict ``parameters``, + # or dict / JSON-string ``arguments``). An ordinary JSON answer like + # {"name":"web_search","result":"no call"} is content, so the strip keeps 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 (e.g. in a "result" value of an + # ordinary JSON answer) is data, and stripping it name-agnostically would delete content. + 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, + *, + in_mapping: bool = False, +): + """Parse one Gemma arg value at ``i`` in a single O(n) forward pass; returns + ``(value, next_index, closed)``. ``closed`` is False when a string/object/array + runs off the end without its terminator, so the caller can fall back to raw. + ``in_mapping`` applies the top-level rule that a comma only ends the value + when a ``key:`` follows (array elements split on every top-level comma).""" + 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), False + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True + if text[i] == "{": + return _gemma_parse_mapping(text, i) + if text[i] == "[": + return _gemma_parse_array(text, i) + if text[i] in "\"'": + # Raw-quoted string: delimiters inside are data (``{city:"New, York"}`` is one + # value); returned unquoted like the top-level scalar coercion. + quote = text[i] + j = i + 1 + n = len(text) + while j < n: + if text[j] == "\\" and j + 1 < n: + j += 2 + continue + if text[j] == quote: + return text[i + 1 : j], j + 1, True + j += 1 + return text[i + 1 :], n, False + # Primitive / unquoted code: same delimiter rules as the top-level scan (bracket depth + # + contextual quote openers hide commas and closers). + end = i + n = len(text) + depth = 0 + quote = "" + prev = ":" + prev_raw = ":" + while end < n and not text.startswith(_GEMMA_STR_BEGIN, end): + ch = text[end] + if quote: + if ch == "\\" and end + 1 < n: + end += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth == 0: + break + depth -= 1 + elif ch == "," and depth == 0: + if not in_mapping or _GEMMA_KEY_RE.match(text, end + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + end += 1 + if end == i: + # Stray delimiter where a value was expected: consume one char so callers always + # advance (no infinite loop on malformed input). + return "", i + 1, True + raw = text[i:end].strip() + if raw == "true": + return True, end, True + if raw == "false": + return False, end, True + if raw == "null": + return None, end, True + try: + return int(raw), end, True + except ValueError: + pass + try: + return float(raw), end, True + except ValueError: + pass + return raw, end, True + + +def _gemma_parse_array(text: str, start: int): + """Parse a Gemma ``[...]`` array at ``text[start] == '['`` in one forward + pass; returns ``(list, next_index, closed)``.""" + items: list[Any] = [] + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "]": + return items, i + 1, True + if i >= n: + break + v, i, _closed = _gemma_parse_value(text, i) + items.append(v) + return items, i, False + + +def _gemma_coerce_scalar(raw: str) -> Any: + """Coerce an unquoted Gemma value to bool/int/float/None, else keep str + (quotes stripped first so quoted/unquoted variants compare identical).""" + raw = raw.strip() + if len(raw) >= 2 and raw[0] == raw[-1] and raw[0] in "\"'": + return raw[1:-1] + if raw == "true": + return True + if raw == "false": + return False + if raw == "null": + return None + try: + return int(raw) + except ValueError: + pass + try: + return float(raw) + except ValueError: + pass + return raw + + +def _gemma_strip_quoted_leaves(value: Any) -> Any: + """Recursively unquote quoted string leaves of a nested stripped-stream value, + so nested ``city:"New York"`` matches the top-level coercion (no stray quotes).""" + if isinstance(value, str): + v = value.strip() + if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'": + return v[1:-1] + return value + if isinstance(value, dict): + return {k: _gemma_strip_quoted_leaves(v) for k, v in value.items()} + if isinstance(value, list): + return [_gemma_strip_quoted_leaves(v) for v in value] + return value + + +def _gemma_parse_stripped_body(body: str) -> dict[str, Any]: + """Parse a quote-less Gemma arg body ``key:value, key2:value2`` (the + ``skip_special_tokens`` stream with ``<|"|>`` markers removed). Each value runs + to the next top-level ``, key:`` boundary, tracking ``{}``/``[]``/``()`` depth so + commas/braces inside a ``code`` / ``command`` value aren't truncated.""" + out: dict[str, Any] = {} + i, n = 0, len(body) + while i < n: + m = _GEMMA_KEY_RE.match(body, i) + if not m: + break + key = m.group(1) + i = m.end() + vstart = i + depth = 0 + quote = "" + # Contextual quote openers mirror _gemma_body_brace_end. + prev = ":" + prev_raw = ":" + while i < n: + ch = body[i] + if quote: + # A ``, key:`` shape inside the quoted string is not a boundary. + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth > 0: + depth -= 1 + elif ch == "," and depth == 0 and _GEMMA_KEY_RE.match(body, i + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + raw_val = body[vstart:i].strip() + if raw_val[:1] in "{[": + # Nested object/array: accept only a fully consumed, closed parse; a + # truncated/malformed value falls back to the raw string. + parsed, end, closed = _gemma_parse_value(raw_val, 0) + out[key] = ( + _gemma_strip_quoted_leaves(parsed) + if (closed and end == len(raw_val)) + else _gemma_coerce_scalar(raw_val) + ) + else: + out[key] = _gemma_coerce_scalar(raw_val) + if i < n and body[i] == ",": + i += 1 + return out + + +def _gemma_parse_mapping(text: str, start: int): + """Parse a Gemma ``{key:value, ...}`` mapping at ``text[start] == '{'`` in one + forward pass; returns ``(dict, next_index, closed)`` (``closed`` True iff the + matching ``}`` was reached).""" + out: dict[str, Any] = {} + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "}": + return out, i + 1, True + if i >= n: + break + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = text[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and text[i] not in ":}": + i += 1 + key = text[kstart:i].strip() + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == ":": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + if text[i] == "}": + out[key] = None + return out, i + 1, True + v, i, _closed = _gemma_parse_value(text, i, in_mapping = True) + out[key] = v + return out, i, False + + +# ── DeepSeek R1 / V3 / V3.1 ───────────────────────────────────────── + + +def _find_outside_json_strings(text: str, needle: str, start: int) -> int: + """Index of ``needle`` at/after ``start`` OUTSIDE any JSON string, or -1: a + marker inside an argument string must not be taken as the structural terminator.""" + i = start + n = len(text) + in_string = False + esc = False + while i < n: + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + i += 1 + continue + if text.startswith(needle, i): + return i + i += 1 + return -1 + + +def _parse_deepseek_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """DeepSeek R1 / V3 / V3.1. + + R1: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`<｜tool▁call▁end｜>...`` + V3.x: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` + + Mirrors llama.cpp's pre-autoparser ``common_chat_parse_deepseek_r1`` / + ``_v3_1`` handling; tolerates the 5 opener variants llama.cpp keeps. + """ + out: list[dict] = [] + begin = _DEEPSEEK_BEGIN_RE.search(content) + if not begin: + return out + scan_start = begin.end() + # Envelope end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would truncate the call. + end_pos = _find_outside_json_strings(content, _DEEPSEEK_END, scan_start) + # Strict mode: an unclosed envelope is truncated; reject, don't heal to EOF. + if not allow_incomplete and end_pos < 0: + return out + scan_end = end_pos if end_pos >= 0 else len(content) + body = content[scan_start:scan_end] + + # R1 path first: ``function<｜tool▁sep｜>NAME\n```json\n{...}\n```<｜tool▁call▁end｜>``. + pos = 0 + while pos < len(body): + fpos = body.find(_DEEPSEEK_R1_FUNC_MARKER, pos) + if fpos < 0: + break + name_start = fpos + len(_DEEPSEEK_R1_FUNC_MARKER) + nl = body.find("\n", name_start) + if nl < 0: + break + if not body.startswith(_DEEPSEEK_R1_FENCE, nl): + pos = name_start + continue + name = body[name_start:nl].strip() + json_start = nl + len(_DEEPSEEK_R1_FENCE) + # Walk a balanced ``{`` even if the trailing fence is truncated. + if json_start >= len(body) or body[json_start] != "{": + pos = json_start + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + # The closing fence + <｜tool▁call▁end｜> must IMMEDIATELY follow the JSON, else an + # unbounded search lands on a LATER call's terminator. Absent close: heal past the + # JSON (strict rejects); later well-formed calls are still kept. + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + close_m = _DEEPSEEK_R1_CLOSE_RE.match(body, after) + if not allow_incomplete and close_m is None: + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close_m.end() if close_m else brace_end + 1 + if out: + return out + + # V3 / V3.1: name then bare JSON. Use ``str.find`` for the sep marker and walk + # back for the name (a ``[^\n<]+`` regex search is O(N^2) on truncated bodies). + pos = 0 + while pos < len(body): + sep_pos = body.find(_DEEPSEEK_SEP, pos) + if sep_pos < 0: + break + # Walk left from sep_pos to the name start; stop at ``\n`` (turn boundary), ``<`` + # (tag start), or ``>`` (end of an optional ``<｜tool▁call▁begin｜>``). + name_start = sep_pos + while name_start > pos and body[name_start - 1] not in "\n<>": + name_start -= 1 + name = body[name_start:sep_pos].strip() + json_start = sep_pos + len(_DEEPSEEK_SEP) + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = sep_pos + len(_DEEPSEEK_SEP) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + # Strict mode: a real V3 call closes with the per-call <｜tool▁call▁end｜>; without + # it the call is truncated/merged, so skip it but keep scanning for a later + # well-formed call (matches Kimi strict). + if not allow_incomplete: + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + if not body.startswith(_DEEPSEEK_CALL_END, after): + pos = brace_end + 1 + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance just past the JSON; seeking the optional <｜tool▁call▁end｜> could land on + # a LATER call's end marker and skip the call between. + pos = brace_end + 1 + return out + + +# ── GLM 4.5 / 4.6 / 4.7 ───────────────────────────────────────────── + + +def _parse_glm_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """GLM 4.5 / 4.6 / 4.7. + + ``NAME[\\n]K[\\n]V + ...``. Multi-call is back-to-back blocks, no envelope. + Mirrors llama.cpp's GLM 4.x tool-call handling (``common_chat_params_init_glm_4_5`` + plus its generalized XML-style parser, llama.cpp PRs #15904 / #16932). + """ + out: list[dict] = [] + pos = 0 + while pos < len(content): + m = _GLM_TC_OPEN_RE.search(content, pos) + if not m: + break + name = m.group(1).strip() + apos = m.end() # absolute position in ``content``; advances past each pair + + args: dict[str, Any] = {} + valid = True + close = -1 + # Walk arg pairs directly against ``content``: a value may contain a literal + # , so the real close is the before the next . + # ``str.find`` keeps this linear. + while True: + ks = content.find(_GLM_ARG_KEY_OPEN, apos) + tc = content.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = content.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < len(content) and content[vstart] in " \t\r\n": + vstart += 1 + if not content.startswith(_GLM_ARG_VAL_OPEN, vstart): + # Key without : strict rejects the call; Auto-Heal skips it. + if not allow_incomplete: + valid = False + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + # A first-match find on would truncate values containing literal + # close tags and execute corrupted arguments. + ve = _glm_value_close(content, vs, strict = not allow_incomplete) + key = content[ks + len(_GLM_ARG_KEY_OPEN) : ke].strip() + if ve < 0: + # Unclosed : strict rejects the whole call; Auto-Heal keeps the + # partial value (a truncated query is not a no-arg call). + if not allow_incomplete: + valid = False + break + # Bound the healed value at the next structural tag, not EOF, so a value + # missing only its can't swallow the markup after it. + nk = content.find(_GLM_ARG_KEY_OPEN, vs) + tc = content.find(_GLM_TC_CLOSE, vs) + bounds = [b for b in (nk, tc) if b >= 0] + if not bounds: + args[key] = content[vs:].rstrip() + break + bound = min(bounds) + args[key] = content[vs:bound].rstrip() + apos = bound + continue + raw_val = content[vs:ve] + apos = ve + len(_GLM_ARG_VAL_CLOSE) + # Decode only unambiguous JSON literals; else keep the value RAW so whitespace + # in string args survives (matches vLLM glm4_moe). ``"`` is left out of the + # probe: a verbatim string's quotes are meaningful. + probe = raw_val.strip() + if ( + probe[:1] in "{[" + or probe in ("true", "false", "null") + or _GLM_JSON_NUMERIC_RE.fullmatch(probe) + ): + try: + args[key] = json.loads(probe) + continue + except (json.JSONDecodeError, ValueError): + pass + args[key] = raw_val + + # Strict mode: a block with no is truncated; reject it. + if not allow_incomplete and close < 0: + valid = False + + if name and valid: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close + len(_GLM_TC_CLOSE) if close >= 0 else len(content) + return out + + +# ── Kimi K2 / Moonshot ────────────────────────────────────────────── + + +def _parse_kimi_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Kimi K2. + + ``<|tool_calls_section_begin|><|tool_call_begin|>functions.NAME:IDX + <|tool_call_argument_begin|>{json}<|tool_call_end|>... + <|tool_calls_section_end|>``. Full id is preserved on ``tool_calls + [i].id`` for round-trip through the chat template. Outer loop walks + every section in the stream (vLLM / SGLang parity); mirrors llama.cpp's + Kimi K2 handling via its generalized XML-style parser (llama.cpp PR #16932). + """ + out: list[dict] = [] + outer_pos = 0 + while True: + section_start = content.find(_KIMI_SECTION_BEGIN, outer_pos) + if section_start < 0: + break + scan_start = section_start + len(_KIMI_SECTION_BEGIN) + # Section end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would drop the later valid call. + section_end = _find_outside_json_strings(content, _KIMI_SECTION_END, scan_start) + scan_end = section_end if section_end >= 0 else len(content) + body = content[scan_start:scan_end] + # Truncated tail: parse what we have, then exit. In strict mode a section with no + # <|tool_calls_section_end|> is truncated; reject it instead. + if section_end < 0: + if allow_incomplete: + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = True + ) + ) + return out + outer_pos = section_end + len(_KIMI_SECTION_END) + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = allow_incomplete + ) + ) + + # The section wrapper is optional (llama.cpp): a bare <|tool_call_begin|> call parses + # as one section when the loop matched nothing. + if not out and _KIMI_CALL_BEGIN in content: + out.extend( + _parse_kimi_section_body( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + ) + return out + + +def _parse_kimi_section_body( + body: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse one Kimi K2 section body (between begin / end markers).""" + out: list[dict] = [] + pos = 0 + while pos < len(body): + call_start = body.find(_KIMI_CALL_BEGIN, pos) + if call_start < 0: + break + id_start = call_start + len(_KIMI_CALL_BEGIN) + arg_begin = body.find(_KIMI_ARG_BEGIN, id_start) + if arg_begin < 0: + break + full_id = body[id_start:arg_begin].strip() + m = _KIMI_ID_RE.match(full_id) + if m: + # group(1) is the whole name; do NOT split on ``.`` -- a dotted MCP name stays intact. + name = m.group(1) + else: + base = full_id.split(":")[0] + name = base[len("functions.") :] if base.startswith("functions.") else base + # Drop bare-counter ids (``3``, ``42``) -- matches vLLM; SGLang infers the name + # from the tool schema, which we don't have here. + if name.isdigit(): + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + brace_end = ( + _balanced_brace_end(body, json_start) + if (json_start < len(body) and body[json_start] == "{") + else None + ) + if brace_end is None: + pos = arg_begin + len(_KIMI_ARG_BEGIN) + else: + pos = brace_end + 1 + continue + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + # Balanced brace lets a truncated trailing end marker still surface a call. + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = arg_begin + len(_KIMI_ARG_BEGIN) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + # Malformed / truncated JSON: skip this call but keep parsing later ones + # instead of dropping the rest of the section (vLLM recovers them). + nxt = body.find(_KIMI_CALL_BEGIN, json_start) + if nxt < 0: + break + pos = nxt + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if not allow_incomplete: + # Strict mode: this call must close with <|tool_call_end|> before the next + # <|tool_call_begin|>; otherwise it is truncated, so reject it. + end_marker = body.find(_KIMI_CALL_END, brace_end + 1) + next_call = body.find(_KIMI_CALL_BEGIN, brace_end + 1) + if end_marker < 0 or (next_call >= 0 and end_marker > next_call): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": full_id or f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance past the JSON; seeking <|tool_call_end|> could skip a following call + # when this one's end marker is missing. + pos = brace_end + 1 + return out diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e8367ad08c..b91403ed57 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -10,40 +10,65 @@ orchestrator, structlog, httpx, or the rest of the studio backend. import json import re -# Pre-compiled patterns for tool XML stripping. The hyphen in the name -# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, -# issue-number) parse alongside the built-ins. +# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs +# strip first so a closed call goes as a unit before any to-EOF sweep reaches +# nested markup; only the final list adds the .*$ EOF sweeps. +_TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) +_TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), - re.compile(r""), - re.compile(r".*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_PAT, + _TC_FUNC_CLOSED_PAT, + _TC_GEMMA_END_PAT, ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), re.compile(r"<\|tool_call>.*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in +# their argument data cannot make the helper truncate the block and its tail. +_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] +# A lazy closed-pair pattern whose close token is absent rescans to EOF from +# every opener (quadratic, re-run per streamed token); skip that doomed pass. +_PAT_REQUIRED_TOKEN = { + _TC_JSON_CLOSED_PAT: "", + _TC_GEMMA_CLOSED_PAT: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """Apply ``patterns`` in order, skipping closed-pair passes with no close token.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + text = pat.sub("", text) + return text + # 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 = "" _FUNC_CLOSE_TAG = "" -# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next -# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. -# `location:New York, NY`), so it must not terminate the value. The key token -# 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*:") +# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next +# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`, +# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") def _balanced_brace_end( @@ -139,14 +164,8 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """Normalise the elements of a Gemma array value so json.loads succeeds. - - Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of - objects (``items:[{path:a}]``) whose keys/values also lack quotes; left - as-is json.loads fails and the whole call is dropped. Bare string elements - are quoted, object and nested-array elements are normalised recursively, and - quoted strings (already normalised from ``<|"|>``), numbers, and JSON - literals are preserved.""" + """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds: + quote bare strings, recurse into objects/arrays, keep quoted/JSON literals.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -154,11 +173,9 @@ def _quote_gemma_array_elements(body: str) -> str: out.append(element) continue if stripped[0] == "{": - # Object element: quote its keys/bare values like a top-level object. out.append(_quote_gemma_object_keys(stripped)) continue if stripped[0] == "[": - # Nested array: normalise its elements too. inner_end = _balanced_bracket_end(stripped, 0) if inner_end == len(stripped) - 1: out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") @@ -223,7 +240,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 @@ -235,15 +254,12 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 - # Gemma may emit bare string values ({unit:celsius}); quote them so - # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. + # Quote bare string values ({unit:celsius}); JSON stays as-is. ws = i while i < len(src) and src[i].isspace(): i += 1 parts.append(src[ws:i]) if i < len(src) and src[i] == "[": - # Array value: quote bare string elements (e.g. labels:[bug,ui]) - # so json.loads succeeds instead of dropping the call. arr_end = _balanced_bracket_end(src, i) if arr_end < 0: parts.append(src[i:]) @@ -253,9 +269,7 @@ def _quote_gemma_object_keys(src: str) -> str: i = arr_end + 1 elif i < len(src) and src[i] not in '"{': v_start = i - # Consume the bare value up to `}` or a comma that starts the - # next key:value pair; a comma inside the value (e.g. - # `New York, NY`) does not terminate it. + # Bare value: up to `}` or a comma that starts the next key:pair. while i < len(src): if src[i] == "}": break @@ -267,7 +281,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 +306,97 @@ 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 _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """Coverage ``[start, end]`` per marker, used to skip markers that are another + call's data. Closes pair to markers via a per-format stack so an inner close + is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with + a paired close cover through it (markers before the close are data); balanced + without one cover only the braces, so a later sibling is still recovered.""" + n = len(content) + brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] + events = [] # (position, order) with order 0 = braces-done, 1 = close marker + for idx, (_start, brace_end, _kind, _m) in enumerate(markers): + if brace_end >= 0: + events.append((brace_end, 0, _kind, idx)) + for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): + for cm in close_re.finditer(content): + # A close inside another call's balanced braces is quoted data; it + # must not pop an earlier close-less marker and swallow a sibling. + if any(s < cm.start() < be for s, be in brace_regions): + continue + events.append((cm.start(), 1, kind, cm.end())) + events.sort(key = lambda e: (e[0], e[1])) + waiting = {"json": [], "gemma": []} + close_end_for: dict[int, int] = {} + for _pos, order, kind, payload in events: + if order == 0: + waiting[kind].append(payload) # marker index, now awaiting its close + elif waiting[kind]: + close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + coverage = [] + for idx, (start, brace_end, _kind, _m) in enumerate(markers): + if brace_end < 0: + coverage.append((start, n)) + elif idx in close_end_for: + coverage.append((start, close_end_for[idx])) + else: + coverage.append((start, brace_end)) + return coverage + + +def _build_markers(content: str): + """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document + order; ``brace_end < 0`` marks an unbalanced (to-EOF) open.""" + markers = [] + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + for m in start_re.finditer(content): + if _inside_open_parameter(content, m.start()): + continue + brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) + markers.append((m.start(), brace_end, kind, m)) + markers.sort(key = lambda c: c[0]) + return markers + + +def marker_coverage(content: str) -> list[tuple[int, int]]: + """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup + inside a marker's coverage (even a marker that failed to parse) as that call's + data rather than a sibling call.""" + return _marker_coverage(content, _build_markers(content)) def parse_tool_calls_from_text( @@ -317,47 +420,40 @@ def parse_tool_calls_from_text( """ tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect every supported call format with spans, then emit in document - # order. A marker inside another call's argument string is data, not a - # separate executable call. - parsed_items = [] # (start, span_end, name, arguments) - candidates = [] # (start, brace_end, kind, match) - for m in _TC_JSON_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1) - if end >= 0: - candidates.append((m.start(), end, "json", m)) - for m in _TC_GEMMA_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True) - if end >= 0: - candidates.append((m.start(), end, "gemma", m)) - candidates.sort(key = lambda c: c[0]) - - candidate_spans = [(s, e) for s, e, _kind, _m in candidates] - for idx, (start, end, kind, m) in enumerate(candidates): - if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx): + # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside + # another call's coverage, or an open value, is data not executed. + markers = _build_markers(content) + coverage = _marker_coverage(content, markers) + parsed_items = [] # (start, span_end, name, arguments) in document order + for idx, (start, brace_end, kind, m) in enumerate(markers): + # A marker starting inside another's coverage is that call's data. The + # end is exclusive so a marker at a close's end is an adjacent sibling. + if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue + if brace_end < 0: + continue # unclosed: not parseable; the fallback still excludes its XML if not allow_incomplete: - tail = content[end + 1 :].lstrip() + tail = content[brace_end + 1 :].lstrip() close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE if close_re.match(tail) is None: continue try: if kind == "json": - obj = json.loads(content[m.end() - 1 : end + 1]) + obj = json.loads(content[m.end() - 1 : brace_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: name = m.group(1) - arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end])) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - span_end = end + 1 + # Span reaches through the close tag when present, else just the braces. + span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) close_m = close_re.match(content, span_end + ws) @@ -365,11 +461,15 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) + # Function-XML calls promote in document order alongside marker calls (the + # #6801 contract). A inside any marker's coverage is excluded -- + # even if that marker failed to parse -- so nested XML cannot escape; one + # after a balanced close-less marker is a sibling, not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in candidate_spans) + and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -382,7 +482,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 +504,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 +522,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,11 +544,107 @@ def parse_tool_calls_from_text( } ) call_spans.append((start, span_end)) + if with_spans: return tool_calls, call_spans return tool_calls +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native spans, brace/quote-balanced so a literal + ```` in a quoted argument cannot truncate the span. An incomplete + span is dropped to EOF when ``final``, else kept (still streaming).""" + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + # Unbalanced: nothing completes from here on. Drop the rest if final, + # else keep it; stop either way (rescanning would be quadratic). + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + # Junk between } and is malformed-call markup: strip through + # the close, keep text after it. No close anywhere means stop (linear). + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + out.append(text[cursor:start]) + cursor = close.end() + out.append(text[cursor:]) + return "".join(out) + + +def _gemma_span_ranges(text: str) -> list: + """``(start, end)`` of each complete Gemma-native span; same walk as + ``_strip_gemma_native_spans`` without stripping.""" + ranges: list[tuple] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + break + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + break + ranges.append((start, close.end())) + cursor = close.end() + return ranges + + +def _strip_closed_blocks_outside_gemma(text: str) -> str: + """Closed JSON/function pre-pass that skips matches starting inside a complete + Gemma span: deleting across the span boundary would mangle the Gemma close and + truncate the tail. A skipped match resumes at the covering span's end, so a + real function-XML call after the span is still stripped.""" + ranges = _gemma_span_ranges(text) + if not ranges: + return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + for pat in _TOOL_CLOSED_BLOCK_PATS: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + out: list[str] = [] + pos = 0 + while True: + m = pat.search(text, pos) + if m is None: + out.append(text[pos:]) + break + covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None) + if covering is not None: + out.append(text[pos : covering[1]]) + pos = covering[1] + continue + out.append(text[pos : m.start()]) + pos = m.end() + new_text = "".join(out) + if new_text != text: + text = new_text + ranges = _gemma_span_ranges(text) + return text + + +def strip_tool_markup_final(text: str) -> str: + """Final display strip, shared with the streaming wrappers so all paths order + the passes identically: Gemma-aware closed JSON/function blocks first, then + well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed + spans and drop any unclosed remainder to EOF. Whitespace is kept.""" + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = True) + return strip_tool_patterns(text, _TOOL_ALL_PATS) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -456,7 +652,9 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. """ - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + if final: + return strip_tool_markup_final(text).strip() + # Non-final: same ordering as the final path, but incomplete blocks are kept. + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = False) + return strip_tool_patterns(text, _TOOL_CLOSED_PATS) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 91a15852a5..8e06b66784 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( @@ -841,17 +852,14 @@ class _SameTaskStreamingResponse(StreamingResponse): **kwargs, ) -> None: super().__init__(*args, **kwargs) - # Async callable invoked when the client disconnects before the body - # iterator is ever advanced. A generator that never started cannot run - # its own try/finally, so a stream that acquires resources before its - # first yield (the passthrough opens an upstream httpx stream eagerly) - # passes this to release them. + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. self._unstarted_cleanup = unstarted_cleanup async def __call__(self, scope, receive, send) -> None: - # Track whether the body iterator was ever advanced: send() only emits a - # body message after the generator yields its first chunk, so a failure - # before then means it never entered its try/finally. + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. body_started = False async def _tracking_send(message) -> None: @@ -862,15 +870,11 @@ class _SameTaskStreamingResponse(StreamingResponse): try: await self.stream_response(_tracking_send) - except OSError: - # Client disconnected mid-send. + except OSError: # client disconnected mid-send if body_started: - # The generator produced at least one chunk and is suspended in - # its try/finally. Throw CancelledError into it (not aclose's - # GeneratorExit) so its `except asyncio.CancelledError` handler - # runs and finishes any api_monitor entry; GeneratorExit would - # skip it and only run `finally`. Fall back to aclose() without - # athrow. + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. athrow = getattr(self.body_iterator, "athrow", None) if athrow is not None: try: @@ -882,16 +886,16 @@ class _SameTaskStreamingResponse(StreamingResponse): if aclose is not None: await aclose() else: - # http.response.start failed before the body iterator advanced, - # so its try/finally never armed and aclose()/athrow() are no-ops - # on an unstarted generator. Release any resources acquired - # before the first yield via the explicit cleanup hook. + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - if self._unstarted_cleanup is not None: + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -899,6 +903,16 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -1148,6 +1162,13 @@ 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_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) +from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS from core.inference.passthrough_healing import ( StreamToolCallHealer, heal_gate, @@ -1306,6 +1327,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), mirroring the parser's tolerance. +_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 @@ -1316,17 +1342,22 @@ 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 any supported parser recognises (template advertises tools but + # uses none -> drop the pill). Reuse the parser's own signal list so this + # gate never drifts (a hand-maintained copy lost the DeepSeek variants); + # ```` is GLM's unique signal, absent from the shared set. The + # bare-JSON ``{"name":`` form is matched below with the whitespace/escape- + # tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped + # templates are not mis-classified as tool-less. + _PARSER_MARKERS = ( + *_PARSER_TOOL_SIGNALS, + "", + ) 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 " @@ -1347,6 +1378,39 @@ 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/Qwen3.5/GLM prefill it). + + Gated on the STANDARD ````/```` markers: a bespoke reasoning channel (e.g. gemma) + never emits ````, so prefilled mode would swallow the whole answer -- excluded, as are + gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill. + """ + 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 + # A reasoning_effort="none" request disables thinking for enable_thinking_effort + # (GLM-5.2) models the same way enable_thinking=False does (see + # ``_request_reasoning_kwargs``). Without this, the model emits no ```` and + # a plain answer is swallowed whole into reasoning_content, leaving the visible + # response empty. + 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. @@ -1617,30 +1681,83 @@ 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 four leak +# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN +# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a +# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation +# is the parser's own, so a signal we parse is never left un-stripped. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + _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)" + # Arm order/notes: the closed ```` arm runs first and extends + # to the call's REAL close so a literal ```` in a value does not + # leak the tail; the combined arm still catches ```` and orphan + # tails. The python_tag arm bounds only on REAL Llama control sentinels + # (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values). + # The last arms cover DeepSeek envelopes (all opener variants), Kimi section + # blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser. + # Those three arms carry a call-shaped lookahead (matching the parser's + # ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker + # (``See <|tool_call_begin|> in the docs``) is only stripped when a real + # call actually follows the marker, or the marker is a bare fragment at EOF. + 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)\|))*" + r"|" + + _DS_OPEN_SRC + + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" + r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)" + r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)" + # ```` is the attribute-form alias of ```` (the parser accepts + # both); strip a tail-only orphan close of either spelling. + r"|\s*\Z", _re.DOTALL, ) -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.""" +def _gemma_strip_gate(tools) -> set: + """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the + parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools + enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is + then prose, and ``None`` would strip-all and delete a legitimate answer.""" + names = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Combine the parser's scan-based strips (Mistral balanced-brace, gated + Gemma wrapper-less, GLM real-close, guarded function-XML) with + ``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so + literal markup inside argument values is data, not a leaked tail. + ``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call.""" + cleaned = _strip_glm_calls( + _strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names), + final = True, + ) + cleaned = _strip_function_xml_calls(cleaned, final = True) + return _TOOL_XML_RE.sub("", cleaned) + + +def _strip_tool_xml_for_display( + text: str, + *, + auto_heal_tool_calls: bool, + enabled_tool_names: Optional[set] = None, +) -> str: + """Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml`` + so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no + ``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip.""" if not auto_heal_tool_calls: return text - return _TOOL_XML_RE.sub("", text) + return _strip_tool_xml(text, enabled_tool_names) logger = get_logger(__name__) @@ -4046,12 +4163,9 @@ async def generate_stream( _DONE = object() while True: if cancel_event.is_set(): - # The disconnect watcher set cancel_event between chunks. - # Reset the backend here: closing the Python generator does - # not signal a subprocess backend, so without this it keeps - # decoding after the client is gone. The finally's reset is - # guarded on cancel_event being unset, so it will not run - # again for this path. + # Watcher set cancel_event between chunks. Reset here: closing + # the generator does not signal a subprocess backend, so it would + # keep decoding. The finally's reset is guarded, so no double-run. backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) @@ -5661,6 +5775,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5936,6 +6051,7 @@ async def openai_chat_completions( _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ).strip() def gguf_generate_with_tools(): @@ -6069,6 +6185,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6140,6 +6257,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6174,6 +6292,7 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ) return full_text, usage, finish finally: @@ -6396,6 +6515,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6545,6 +6665,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()) @@ -6631,6 +6767,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ).strip(), } ) @@ -6686,6 +6823,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 (GGUF parity); 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() @@ -6702,7 +6852,11 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: + # Iteration boundary: flush reasoning, then start a fresh extractor for the next turn. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() status_data = json.dumps( { "type": "tool_status", @@ -6714,7 +6868,11 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + # Flush reasoning before the tool_start line 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 @@ -6723,14 +6881,24 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ) new_text = clean_cumulative[len(prev_text) :] 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 @@ -6783,6 +6951,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6804,22 +6973,32 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - api_monitor.set_reply(monitor_id, content_text) + # Split prefilled reasoning out of the visible answer (GGUF parity); 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", ) ], @@ -6898,6 +7077,8 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" + # Split prefilled into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves 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 @@ -6926,9 +7107,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. @@ -6975,6 +7168,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6990,18 +7184,27 @@ async def openai_chat_completions( for token in generate(): full_text = token + # Split prefilled reasoning from the visible answer (GGUF parity); 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")) @@ -7824,10 +8027,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 ````. Callers default False. + 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, @@ -7850,14 +8061,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 @@ -7894,7 +8112,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, "") @@ -7903,8 +8121,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 @@ -9734,7 +9956,9 @@ 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"], _gemma_strip_gate(openai_tools) + ).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -9779,6 +10003,7 @@ async def anthropic_messages( message_id, model_name, disable_parallel_tool_use = _disable_parallel, + openai_tools = openai_tools, ) ) @@ -9857,11 +10082,8 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between events, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # event or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # events, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -9888,7 +10110,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"], _gemma_strip_gate(openai_tools)) # 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). @@ -9953,11 +10175,8 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between chunks, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # chunk or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # chunks, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -10046,6 +10265,7 @@ async def _anthropic_tool_non_streaming( message_id, model_name, disable_parallel_tool_use = False, + openai_tools = None, ): """Non-streaming response for the tool-calling path. @@ -10074,7 +10294,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"], _gemma_strip_gate(openai_tools)) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10544,9 +10764,13 @@ async def _anthropic_passthrough_non_streaming( 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. + # only for opted-out or no-client-tool requests. Use the full + # _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded + # function-XML leaks are cleaned too, not just _TOOL_XML_RE forms, + # with the Gemma display gate so a disabled/example call:NAME{...} + # in prose survives. if not healing_active: - text = _TOOL_XML_RE.sub("", text) + text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools)) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) @@ -10909,6 +11133,10 @@ async def _openai_passthrough_stream( response ``id``, ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so the client sees a standard OpenAI response. + + Reasoning/tool-call splitting is delegated to llama-server (``--jinja + --reasoning-format auto``), so ``delta.content`` carries no raw markup and is + deliberately not re-parsed locally, unlike the ``/completion`` paths. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( @@ -11325,11 +11553,9 @@ async def _openai_passthrough_stream( delta = choice.get("delta") if isinstance(delta, dict) and delta.get("tool_calls"): saw_tool_call_delta = True - # Detect an upstream error chunk independently of API - # monitoring: when monitor_id is None (skip_api_monitor), - # _monitor_openai_sse_line returns before inspecting the - # error, so without this the synthetic-finish guard would - # emit a successful finish_reason after a failed stream. + # Detect an error chunk independently of API monitoring + # (skip_api_monitor returns early), else the synthetic + # finish would fire after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True # With healing active, a content-bearing line may be replaced by diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py new file mode 100644 index 0000000000..75d0117015 --- /dev/null +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Mapper models whose own tokenizer ships no chat_template have their turn-end +eos resolved at LOAD from an empty template (document eos only). The effective +template is installed later, at generate time, via get_chat_template, so the +turn-end-eos cache must be refreshed then; otherwise generate_stream runs past +the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes). +""" + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +# These tests construct InferenceBackend, pulling the full stack. CI may lack +# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup +# (RuntimeError); skip at module level so collection is not aborted (exit 2). +try: + from core.inference import inference as inf_mod # noqa: E402 + from core.inference.inference import InferenceBackend # noqa: E402 +except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent + pytest.skip( + f"full inference backend unavailable ({type(exc).__name__}: {exc})", + allow_module_level = True, + ) + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_GEMMA = "{% for m in messages %}{{m.role}}\n{{m.content}}{% endfor %}" + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.pad_token_id = eos_id + self.unk_token_id = None + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok) + + +def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/qwen2.5-0.5b" + + # No chat_template at load, so the cache stored only the document eos, though + # <|im_end|> is atomic in the vocab (unused until the mapper installs a template). + bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645}) + model_info = { + "tokenizer": bare_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [151643], + } + backend.models = {backend.active_model_name: model_info} + + # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time. + templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False + ) + + # Stub the tail so the generator runs through the refresh without a real model. + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # After the template is applied the cache must include the ChatML turn-end id. + assert model_info["chat_turn_end_eos_ids"] == [151643, 151645] + + +def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch): + # Regression: get_chat_template can return a remapped tokenizer (Gemma: + # folded onto the eos id) while generate_stream re-reads the original. Resolving on + # the swap yields a narrower set, so the refresh must UNION, never overwrite. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-2b-it" + + # Original tokenizer (used by generate_stream): =107 distinct from + # eos=1, so the load-time cache resolved to [1, 107]. + orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 107}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [1, 107], + } + backend.models = {backend.active_model_name: model_info} + + # Destructively-swapped tokenizer: now maps onto eos id 1, so + # resolving on it yields only [1] (drops 107). + swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 1}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The load-time =107 must survive: overwriting with the swapped + # [1] would regress and loop past the turn. + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + + +def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch): + # Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but + # get_chat_template folds it onto the doc-eos id. generate_stream uses the original, + # so read marker strings from the mapped template but ids from the original. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "01-ai/yi-6b" + + # Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7. + orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [2], + } + backend.models = {backend.active_model_name: model_info} + + # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2. + remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2. + assert model_info["chat_turn_end_eos_ids"] == [2, 7] + + +class _FakeProcessor: + """A ProcessorMixin-like container: carries the chat_template itself and + wraps the real text tokenizer as ``.tokenizer`` (the vision layout).""" + + def __init__(self, chat_template, tokenizer): + self.chat_template = chat_template + self.tokenizer = tokenizer + + +def test_resolve_chat_eos_reads_vision_processor_template(): + # Vision model: the chat_template lives on the processor while the inner tokenizer + # ships none. _resolve_chat_eos must read the marker from the processor but resolve + # its id on the inner tokenizer, and repair generation_config. + from types import SimpleNamespace + + inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"": 107}) + processor = _FakeProcessor(_GEMMA, inner_tok) + model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1)) + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-3-4b-it" + model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True} + backend.models = {backend.active_model_name: model_info} + + backend._resolve_chat_eos(backend.active_model_name) + + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + # generation_config repaired so the vision .generate() path stops at the turn. + assert model.generation_config.eos_token_id == [1, 107] diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py new file mode 100644 index 0000000000..13d1ecabaa --- /dev/null +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""apply_chat_template_for_generation must coerce assistant tool_call arguments +from the OpenAI JSON-string form to a dict before rendering. Strict tool +templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and +raise "Can only get item pairs from a mapping." on the string form when a prior +tool call is re-rendered on the next turn (MLX + transformers paths). +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_template_helpers import ( # noqa: E402 + _normalize_tool_call_arguments, + apply_chat_template_for_generation, +) + + +def _conv(arguments): + return [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "id": "c1", + "function": {"name": "web_search", "arguments": arguments}, + } + ], + }, + {"role": "tool", "name": "web_search", "content": "21C sunny"}, + ] + + +class _StrictTemplateTokenizer: + """Mimics a strict Qwen tool template: rejects string tool_call arguments.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise TypeError("Can only get item pairs from a mapping.") + return "RENDERED" + + +def test_string_arguments_are_parsed_to_dict(): + out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}')) + args = out[1]["tool_calls"][0]["function"]["arguments"] + assert args == {"query": "sweden"} + + +def test_dict_arguments_untouched_and_no_copy(): + conv = _conv({"query": "sweden"}) + assert _normalize_tool_call_arguments(conv) is conv + + +def test_non_json_string_left_as_is(): + out = _normalize_tool_call_arguments(_conv("not json")) + assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json" + + +def test_render_succeeds_on_strict_template_with_string_arguments(): + # Regression: strict template + string args used to raise. + result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}')) + assert result == "RENDERED" + + +class _RecordingTokenizer: + """Lenient template: renders whatever arguments it is given (string or dict).""" + + def __init__(self): + self.seen_arguments = None + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + self.seen_arguments = call.get("function", {}).get("arguments") + return "RENDERED" + + +def test_lenient_template_receives_original_string_untouched(): + # Lenient template must see the exact original string, not a coerced dict. + tok = _RecordingTokenizer() + apply_chat_template_for_generation(tok, _conv('{"query": "x"}')) + assert tok.seen_arguments == '{"query": "x"}' + + +def test_messages_without_tool_calls_pass_through_unchanged(): + conv = [{"role": "user", "content": "hi"}] + assert _normalize_tool_call_arguments(conv) is conv + + +class _RaiseExceptionTemplateTokenizer: + """Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via + ``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise ValueError( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string." + ) + return "RENDERED" + + +def test_render_succeeds_on_raise_exception_template_with_string_arguments(): + # Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce. + result = apply_chat_template_for_generation( + _RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}') + ) + assert result == "RENDERED" + + +def test_unrelated_template_error_still_propagates_with_dict_args(): + # Failure unrelated to string args (dict args, nothing to coerce) must propagate. + class _AlwaysRaises: + def apply_chat_template(self, messages, **kw): + raise ValueError("template is broken") + + with pytest.raises(ValueError, match = "broken"): + apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"})) diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py new file mode 100644 index 0000000000..c49e39f8fe --- /dev/null +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and +repair generation_config so a chat model whose eos is a bare document terminator +(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn +boundary instead of running past it and looping. Dependency-light: imported here +without the full inference stack. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_eos import ( # noqa: E402 + chat_eos_repair, + resolve_chat_turn_end_eos_ids, + resolve_chat_turn_end_eos_ids_using, +) + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + unk_token_id = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.unk_token_id = unk_token_id + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok, self.unk_token_id) + + +# ---- resolve_chat_turn_end_eos_ids --------------------------------------- + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" + + +def test_qwen35_adds_im_end_from_template(): + # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046). + tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046] + + +def test_marker_in_vocab_but_not_in_template_is_ignored(): + # Base/coder model: <|im_end|> is in the vocab but the template does not use + # it, so it must not become a stop token. + tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044] + + +def test_harmony_template_is_left_untouched(): + # gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end. + harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>" + tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +def test_llama3_eot_id_from_template(): + tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009}) + assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009] + + +def test_gemma4_turn_marker_from_template(): + # Gemma-4 ends turns with while keeping a document eos, so must + # be added as a stop token. + tok = _FakeTokenizer( + 1, chat_template = ".........", token_ids = {"": 106} + ) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 106] + + +def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer(): + # map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id, + # but the original keeps it atomic. Reading marker STRINGS from the template but + # IDS on the original recovers the real turn-end id (7), not the doc-eos id (2). + template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7] + # Same tokenizer for both reproduces the plain resolve (load-time behaviour). + assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2] + + +def test_list_eos_preserved(): + tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 2] + + +def test_missing_marker_maps_to_unk_and_is_skipped(): + tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0) + assert resolve_chat_turn_end_eos_ids(tok) == [7] + + +def test_starling_barred_end_of_turn_from_template(): + # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from + # Gemma's ). eos synced to =2, turn marker at 32000. + starling = "GPT4 Correct Assistant: hi<|end_of_turn|>" + tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000] + + +def test_dict_chat_template_scans_all_variants(): + # Hermes-3 style: chat_template is a {name: template} dict. Detection must scan + # every variant, not bail because the container is not a plain str. + tmpl = {"default": "{{ messages }}", "tool_use": _CHATML} + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_list_of_dicts_chat_template_scans_all_variants(): + # tokenizer_config.json stores multi-templates as a list of {name, template}. + tmpl = [{"name": "default", "template": _CHATML}] + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_dict_harmony_template_left_untouched(): + # A multi-variant container whose variant is harmony must still be left alone. + tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"} + tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +# ---- chat_eos_repair ------------------------------------------------------ + + +def test_repair_adds_missing_turn_end(): + assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046] + + +def test_repair_from_missing_generation_config_eos(): + assert chat_eos_repair(None, [248046]) == [248046] + + +def test_repair_noop_when_already_covered(): + assert chat_eos_repair([248046, 248044], [248046]) is None + + +def test_repair_noop_when_no_turn_end_ids(): + assert chat_eos_repair(248044, []) is None 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..7b653f47aa 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -1,15 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Edge cases in Gemma-native tool-call parsing. - -Covers two failure modes: - 1. A bare (unquoted) string argument that contains a comma, e.g. - ``location:New York, NY`` -- the comma must not be treated as the next - key boundary, or the whole call is dropped. - 2. A tool-call marker that appears INSIDE another call's argument string is - data, not a real call, so it must not be promoted to a second tool call. -""" +"""Gemma-native tool-call parsing edge cases: commas inside bare string values, +and markers inside another call's argument data staying data.""" from __future__ import annotations @@ -21,7 +14,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.tool_call_parser import parse_tool_calls_from_text +from core.inference.tool_call_parser import ( + _gemma_parse_value, + parse_tool_calls_from_text, +) +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -40,14 +37,22 @@ def test_bare_string_argument_with_comma_is_kept(): def test_normal_multi_key_arguments_still_split(): calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') assert len(calls) == 1, calls - # Numbers stay numeric, bare strings get quoted, an explicit quoted comma - # stays inside its value. 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 and dropped the call). + 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. + # A comma before digits-then-colon (timestamp/ratio) is value text, not a key. calls = parse_tool_calls_from_text( "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) @@ -55,9 +60,16 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} +def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept(): + # The wrapper-less Gemma form (no <|tool_call> markers) goes through the + # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE. + calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}") + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"} + + def test_marker_inside_json_argument_is_not_a_second_call(): - # A python call whose `code` argument contains a Gemma marker string. The - # marker is data and must not execute as a second `terminal` call. content = ( '{"name":"python","arguments":{"code":' '"x = 1 # <|tool_call>call:terminal{command:ls}"}}' @@ -75,8 +87,6 @@ def test_two_separate_gemma_calls_both_parse(): def test_mixed_format_calls_preserve_document_order(): - # A Gemma-native call precedes a JSON-format call in the text; tools execute - # in returned order, so `create` must come before `read`. content = ( "<|tool_call>call:create{path:a} then " '{"name":"read","arguments":{"path":"a"}}' @@ -86,8 +96,6 @@ def test_mixed_format_calls_preserve_document_order(): def test_json_marker_inside_gemma_argument_is_not_a_second_call(): - # The reverse of the JSON-outer case: a JSON-style marker inside a Gemma - # call's quoted argument is code text, not a second `terminal` call. content = ( '<|tool_call>call:python{code:<|"|>' 'print({"name":"terminal","arguments":{"command":"ls"}})' @@ -98,18 +106,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): - # An UNQUOTED Gemma value containing a literal marker: the outer object fails - # to normalize (the inner braces/marker break the JSON), but the inner marker - # is nested in the outer candidate span, so it must not be promoted to a - # standalone `terminal` call. The safe outcome is no executed tool call. + # An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer + # candidate span, so it must not be promoted to a standalone `terminal` call (no tool call). content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" calls = parse_tool_calls_from_text(content) assert "terminal" not in [c["function"]["name"] for c in calls], calls def test_bare_string_array_argument_is_quoted(): - # Gemma may emit an array of bare strings without per-element quotes; they - # must be quoted so the call is not dropped. calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") assert len(calls) == 1, calls assert _args(calls[0]) == {"labels": ["bug", "ui"]} @@ -123,8 +127,6 @@ def test_array_keeps_numbers_and_quoted_elements(): def test_array_of_objects_is_normalised(): - # Arrays of objects are a common tool-schema shape; their (unquoted) keys and - # bare values must be normalised too, not left verbatim, or the call drops. calls = parse_tool_calls_from_text( "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" ) @@ -138,9 +140,6 @@ def test_nested_array_elements_are_normalised(): def test_gemma_marker_inside_xml_parameter_is_not_a_second_call(): - # An XML-style call whose value contains a - # Gemma marker: the marker is the parameter's data, not a separate terminal - # call, so only the python call must be returned. content = ( "" "x = 1 # <|tool_call>call:terminal{command:ls}" @@ -159,3 +158,240 @@ 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_wrapperless_nested_object_argument_is_parsed(): + # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. + calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") + assert len(calls) == 1 + assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3} + + +def test_wrapperless_array_argument_is_parsed(): + calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}") + assert len(calls) == 1 + assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2} + + +def test_wrapperless_deeply_nested_object_and_array_are_preserved(): + # The single-pass parser must keep multi-level nesting (objects inside + # objects, arrays inside arrays) intact, not flatten or drop it. + calls = parse_tool_calls_from_text( + "call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}" + ) + assert len(calls) == 1 + assert _args(calls[0]) == { + "loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}}, + "tags": ["a", "b", ["c", "d"]], + "n": 3, + } + + +def test_gemma_parse_array_advances_on_stray_brace(): + # Regression: a stray '}' / ']' / ',' where an array element is expected must + # not stall _gemma_parse_value at the same index (it looped forever before). + from core.inference.tool_call_parser import _gemma_parse_array + + items, end, closed = _gemma_parse_array("[a,}]", 0) + assert end == 5 and closed is True # consumed through the closing ']' + assert items[0] == "a" + + +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # Parse keeps the quoted close marker as data; strip removes the whole span. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" + + +def test_nested_xml_in_malformed_gemma_call_does_not_execute(): + # The failed Gemma candidate's span still covers its nested . + text = ( + "<|tool_call>call:outer{code:id" + ", broken:{x}}" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_unbalanced_gemma_call_with_xml_does_not_execute(): + # Unclosed braces cover to EOF, so the trailing is excluded. + text = ( + "<|tool_call>call:outer{code:" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_standalone_function_xml_still_parses(): + text = "id" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_xml_between_braces_and_close_marker_does_not_execute(): + # Coverage runs to the close marker, so in the gap is data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): + text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_preserves_text_after_malformed_gemma_close(): + # Junk before the close is a malformed span: strip through it, keep the tail. + text = "pre <|tool_call>call:t{a:1} note post" + assert strip_tool_call_markup(text) == "pre post" + assert strip_tool_call_markup(text, final = True) == "pre post" + + +def test_malformed_closed_gemma_span_is_stripped(): + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) + + +def test_valid_call_after_missing_close_is_recovered(): + # A close-less call covers only its braces, so the later call is recovered. + text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" + names_inc = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + ] + assert "b" in names_inc, names_inc + names_strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert names_strict == ["b"], names_strict + + +def test_strip_non_final_keeps_incomplete_gemma_block(): + text = "before <|tool_call>call:t{" + assert strip_tool_call_markup(text) == text + assert strip_tool_call_markup(text, final = True) == "before" + + +def test_json_call_between_gemma_braces_and_close_does_not_execute(): + # A JSON call between the outer's braces and its close is covered data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + '{"name":"terminal","arguments":{"command":"id"}}' + "" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): + # Same escape with a Gemma-native inner marker. + text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): + # The to-EOF Gemma sweep must not eat visible text after . + text = ( + 'before print("<|tool_call>") after' + ) + assert strip_tool_call_markup(text, final = True) == "before after" + assert strip_tool_call_markup(text) == "before after" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # A call-form Gemma opener quoted in a closed block must not truncate it. + xml = "<|tool_call>call:t{" + json_block = ( + '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + ) + for block in (xml, json_block): + text = "before " + block + " after" + assert strip_tool_call_markup(text, final = True) == "before after", block + assert strip_tool_call_markup(text) == "before after", block + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # The close-less marker covers only its braces; the XML sibling is recovered. + text = ( + "<|tool_call>call:bad{broken:{x}} " + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): + # A close token quoted in the later call must not extend the earlier + # close-less marker's coverage over that call. + gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' + names = [ + c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + ] + assert names == ["b"], names + json_text = ( + '{"name":"a","arguments":{}} ' + '{"name":"b","arguments":{"x":""}}' + ) + names_j = [ + c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + ] + assert "b" in names_j, names_j + + +def test_gemma_parse_value_always_advances_on_stray_delimiter(): + # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the + # index by at least one, or a caller looping on it spins forever at 100% CPU (DoS). + for delim in (",", "}", "]"): + text = delim + "rest" + value, nxt, _explicit = _gemma_parse_value(text, 0) + assert nxt > 0, (delim, value, nxt) + + +def test_malformed_gemma_array_does_not_hang(): + # ``[},]`` puts a stray ``}`` at the primitive position inside a list body. + # On the buggy parser this hangs the server; guard with a wall-clock timeout + # so the regression fails 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_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 9e40fbf508..4d73213d15 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP -transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. -CPU-only, no network, no real subprocess (the per-attempt download seam is -monkeypatched). +"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. + +The transport-policy matrix is tested once in unsloth_zoo; here we assert only the +Studio seam: re-exporting the shared API and injecting the marker-aware +prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ from __future__ import annotations -import subprocess import sys -import threading -import time import types as _types from pathlib import Path @@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable deps before importing the module under test. Use the -# real structlog when present; a bare stub left in sys.modules would break later -# modules that log at import time. +# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present; +# a bare stub would break later modules that log at import time. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -34,171 +31,59 @@ except ImportError: sys.modules["structlog"] = _types.ModuleType("structlog") import huggingface_hub -from huggingface_hub import constants as hf_constants + +try: + import unsloth_zoo.hf_xet_fallback as _shared_mod + shared = _shared_mod +except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable + shared = None import utils.hf_xet_fallback as xf -# --------------------------------------------------------------------------- # -# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. -# --------------------------------------------------------------------------- # -REPO = "ztest/xet-watchdog" - - -@pytest.fixture -def hf_cache(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - return tmp_path - - -def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: - d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" - d.mkdir(parents = True, exist_ok = True) - return d - - -def _wait( - predicate, - timeout: float = 2.0, - step: float = 0.02, -) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(step) - return predicate() - - -def test_constant_incomplete_fires_stall(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - assert _wait( - lambda: len(calls) >= 1, timeout = 3.0 - ), "watchdog never fired on a constant-size .incomplete" - finally: - stop.set() - assert "stalled" in calls[0].lower() - - -def test_growing_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - part = blobs / "growing.incomplete" - part.write_bytes(b"\0" * 1024) - - grow_stop = threading.Event() - - def _grow(): - size = 1024 - while not grow_stop.wait(0.05): - size += 4096 - part.write_bytes(b"\0" * size) - - grower = threading.Thread(target = _grow, daemon = True) - grower.start() - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(1.0) # well past stall_timeout, but bytes keep growing - assert calls == [], "watchdog fired despite continuous progress" - finally: - stop.set() - grow_stop.set() - - -def test_no_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(0.8) - assert calls == [], "watchdog fired with no active .incomplete" - finally: - stop.set() - - -def test_stall_fires_at_most_once(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 - ) - try: - assert _wait(lambda: len(calls) >= 1, timeout = 3.0) - time.sleep(0.6) # keep ticking; must not fire again - assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" - finally: - stop.set() - - -def test_get_state_empty_cache(hf_cache): - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_absent_cache_root(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_skips_local_paths(hf_cache): - # Filesystem paths are not HF repo IDs and must be ignored without error. - assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) - - -def test_get_state_sparse_aware(hf_cache): - blobs = _blobs_dir(hf_cache) - sparse = blobs / "sparse.incomplete" - with open(sparse, "wb") as f: - f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks - st = sparse.stat() - if getattr(st, "st_blocks", 0) == 0: - pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") - total, has_incomplete = xf.get_hf_download_state([REPO]) - assert has_incomplete is True - assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" - - -# --------------------------------------------------------------------------- # -# Transport policy: cached short-circuit, cancel, error propagation, and the -# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. -# --------------------------------------------------------------------------- # DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" -@pytest.fixture(autouse = True) -def _no_real_cache_hit(monkeypatch): - """Default: the cached probe misses; tests override it to force a hit.""" +def _requires_shared(): + if shared is None: + pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment") + + +def test_shim_reexports_shared_api(): + _requires_shared() + assert xf.DownloadStallError is shared.DownloadStallError + for name in ( + "start_watchdog", + "get_hf_download_state", + "child_should_disable_xet", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", + ): + assert hasattr(xf, name), f"shim missing {name}" + + +def test_child_should_disable_xet_truth_table(): + assert xf.child_should_disable_xet({"disable_xet": True}) is True + assert xf.child_should_disable_xet({"disable_xet": False}) is False + assert xf.child_should_disable_xet({}) is False + + +def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): + """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + ``prepare_cache_for_transport(..., 'http')`` before the retry.""" + _requires_shared() + for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + seen_disable_xet = [] -class _FakeAttempt: - """Records calls to the download seam and returns scripted results.""" - - def __init__(self, results): - self._results = list(results) - self.calls = [] - - def __call__( - self, + def fake_attempt( repo_id, - filename, - token, *, + kind, + params, + token, repo_type, disable_xet, cancel_event, @@ -208,146 +93,243 @@ class _FakeAttempt: on_status, force_download = False, ): - self.calls.append( - _types.SimpleNamespace( - repo_id = repo_id, - filename = filename, - disable_xet = disable_xet, - repo_type = repo_type, - ) - ) - return self._results[len(self.calls) - 1] + seen_disable_xet.append(disable_xet) + return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None) + monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt) -def _install(monkeypatch, results): - fake = _FakeAttempt(results) - monkeypatch.setattr(xf, "_run_download_attempt", fake) - return fake - - -def test_cached_file_short_circuits(monkeypatch, tmp_path): - cached = tmp_path / "cached.gguf" - cached.write_bytes(b"\0" * 8) - monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) - fake = _install(monkeypatch, []) # must not be called - - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == str(cached) - assert fake.calls == [], "spawned a download for an already-cached file" - - -def test_cancel_before_start_raises_no_attempt(monkeypatch): - fake = _install(monkeypatch, []) - ev = threading.Event() - ev.set() - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) - assert fake.calls == [] - - -def test_nonstall_error_propagates_without_fallback(monkeypatch): - fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) - with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" - assert fake.calls[0].disable_xet is False - - -def test_immediate_success_uses_xet_only(monkeypatch): - prepared = [] - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", - lambda *a, **k: prepared.append(a), - ) - fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == "/cache/model.gguf" - assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False - assert prepared == [], "no cache prep should run when Xet succeeds first try" - - -def test_stall_then_http_fallback_succeeds(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), ) - fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" - assert len(fake.calls) == 2 - assert fake.calls[0].disable_xet is False # Xet first - assert fake.calls[1].disable_xet is True # HTTP fallback - assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + assert seen_disable_xet == [False, True] # Xet first, then HTTP + assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" -def test_second_stall_raises_download_stall_error(monkeypatch): - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("stall", None), ("stall", None)]) - with pytest.raises(xf.DownloadStallError): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 2 +def test_shim_snapshot_injects_studio_prepare(monkeypatch): + """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + captured = {} + + def fake_snapshot(repo_id, **kwargs): + captured["repo_id"] = repo_id + captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") + return "/tmp/snap-dir" + + monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + out = xf.snapshot_download_with_xet_fallback("org/model") + assert out == "/tmp/snap-dir" + assert captured["repo_id"] == "org/model" + assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http -def test_cancelled_midattempt_raises_no_fallback(monkeypatch): - fake = _install(monkeypatch, [("cancelled", None)]) - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1 +def test_degrades_gracefully_without_shared_helper(monkeypatch): + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" + import importlib + + class _BlockShared: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + raise ModuleNotFoundError(f"No module named '{name}'", name = name) + return None + + finder = _BlockShared() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + + # Boots without raising and mirrors the shared API surface. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.child_should_disable_xet({"disable_xet": True}) is True + assert degraded.get_hf_download_state(["x"]) is None # unmeasurable + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() # never fires + + # Degraded mode still emits heartbeats so the inactivity deadline is not tripped. + import time as _time + + beats = [] + hb_stop = degraded.start_watchdog( + repo_ids = ["x"], + on_stall = lambda m: None, + on_heartbeat = beats.append, + interval = 0.02, + ) + try: + deadline = _time.monotonic() + 2.0 + while not beats and _time.monotonic() < deadline: + _time.sleep(0.02) + assert beats, "degraded watchdog emitted no heartbeat" + finally: + hb_stop.set() + + # Downloads fall back to plain huggingface_hub (no watchdog, no crash). + called = {} + + def _fake_snapshot(repo_id, **kwargs): + called["repo_id"] = repo_id + return "/snap-dir" + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot) + assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir" + assert called["repo_id"] == "org/model" + + # Cancellation still holds: an already-set cancel_event aborts before the HF download. + import threading as _threading + + cancelled = _threading.Event() + cancelled.set() + called.clear() + with pytest.raises(RuntimeError, match = "Cancelled"): + degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + assert "repo_id" not in called, "degraded download ran despite cancellation" + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_per_file_independent_fallback(monkeypatch): - """A stalled shard falls back; a sibling shard that succeeds does not.""" - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" - assert [c.disable_xet for c in fake.calls] == [False, False, True] +def test_degrades_when_unsloth_zoo_entirely_absent(): + """When unsloth_zoo is absent entirely, the import raises + ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still + degrades and does not re-raise, breaking every Studio import that pulls it in.""" + import importlib + + class _BlockZoo: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. + if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): + raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + return None + + finder = _BlockZoo() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Boots without raising and exposes the stub API. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -# --------------------------------------------------------------------------- # -# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect -# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). -# --------------------------------------------------------------------------- # -def _safe_path() -> str: +def test_degrades_when_shared_helper_import_raises_importerror(): + """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only + Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + import importlib + + class _BlockWithImportError: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + # Mirror a torch-less install: a plain ImportError with no .name. + raise ImportError("Unsloth: Pytorch is not installed.") + return None + + finder = _BlockWithImportError() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_zoo = sys.modules.pop("unsloth_zoo", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_zoo is not None: + sys.modules["unsloth_zoo"] = saved_zoo + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): + """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + import importlib import os - return os.environ.get("PATH", "") + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + seen_env = [] -def test_disable_xet_constant_set_in_fresh_interpreter(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " - f"(rc={proc.returncode}): {proc.stderr}" - ) + class _GpuGatedBlocker: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent). + if name == "unsloth_zoo": + # Record the env each attempt sees; raise the no-GPU error both times so the shim + # degrades. + seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise NotImplementedError("Unsloth cannot find any torch accelerator") + return None - -def test_default_leaves_xet_enabled(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"without the env var, constants.HF_HUB_DISABLE_XET was not False " - f"(rc={proc.returncode}): {proc.stderr}" - ) + finder = _GpuGatedBlocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # First attempt without the light env, then a retry with it set. + assert seen_env == [None, "1"], seen_env + # Both attempts raised -> Studio still boots in degraded mode. + assert issubclass(degraded.DownloadStallError, RuntimeError) + # The env override must not leak past the import. + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 05d2a0b80a..dcc759a210 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; derive the 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): + # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the + # shared parser patterns (which know ``[TOOL_CALLS]``); 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,255 @@ 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, + ) + ) + + # The tool ran with the parsed arguments. + 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_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch): + """A truncated markerless object whose "name" is NOT an enabled tool (a person + record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The + end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment + to DRAINING (dropped); it is now gated on the enabled tool names so only a real + truncated tool call is suppressed, ordinary JSON streams through.""" + + truncated = '{"name": "Alice", "age": 30, "bio": "loves ' + 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": "start 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_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 test_gemma_wrapperless_call_streamed_is_not_leaked_and_executes(monkeypatch): + """Gemma 4 GGUF (skip_special_tokens) streams a wrapper-less ``call:NAME{..}`` + with no XML signal. Like bare JSON, the BUFFERING scan must recognise it via + _GEMMA_BARE_TC_RE, drain it silently, and execute the tool -- never leaking + the ``call:`` markup to the user-visible stream.""" + + gemma_call = 'call:web_search{query:"weather in Sydney"}' + first_stream = _streamed_content(gemma_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"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("call:" not in t for t in content_texts), content_texts + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + 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 +2128,191 @@ 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_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): + """A single textual-fallback turn that parses many DISTINCT tool calls must be + capped at _MAX_TOOL_CALLS_PER_TURN (structured delta.tool_calls are grammar + bounded by llama-server; text parsed from content is not). Mirrors the + safetensors loop so one runaway turn cannot fan out into dozens of executions.""" + from core.inference.llama_cpp import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + blocks = "".join( + '{"name":"t%d","arguments":{"i":%d}}' % (i, i) for i in range(n) + ) + first_stream = [_sse({"content": blocks}), _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"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == _MAX_TOOL_CALLS_PER_TURN, [c[0] for c in calls] + # The cap keeps the first calls in order (no reordering / drop of leading ones). + assert [c[0] for c in calls] == [f"t{i}" for i in range(_MAX_TOOL_CALLS_PER_TURN)] + + +def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): + """Exact-duplicate textual calls in one turn collapse to a single execution.""" + blocks = '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + first_stream = [_sse({"content": blocks}), _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"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1, [c[0] for c in calls] + + +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: if re-prompt slots leaked into the budget (the bug) the + # loop would run 2+3=5 rounds; honouring it 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_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 12239e7113..6d26d075cf 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names(): import re as _re from pathlib import Path + from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + src = (Path(__file__).resolve().parent.parent / "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" - ns: dict = {"_re": _re} + ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC} exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns) rx = ns["_TOOL_XML_RE"] stripped = rx.sub( diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 9871965ce8..ac4088fb25 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -100,6 +100,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): ] assert backend._is_vlm is False assert isinstance(backend._tokenizer, _DummyTokenizer) + # Non-LoRA text model: no base_model on the record. + assert backend.models["fake/text"]["base_model"] is None + + +def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch): + # A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template + # lives on the base model. + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace( + identifier = "fake/text-adapter", + is_vision = False, + is_lora = True, + base_model = "fake/text-base", + ) + + assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token") + + record = backend.models["fake/text-adapter"] + assert record["is_lora"] is True + assert record["base_model"] == "fake/text-base" def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( @@ -188,12 +214,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): _install_fake_mlx(monkeypatch) from core.inference.mlx_inference import MLXInferenceBackend - captured = {} + # The text path renders once with tools, then the native-template fallback makes a second no- + # tools probe call (tools=None) to detect whether the template dropped the schema. + captured_calls = [] def _fake_apply(tokenizer, messages, **kwargs): - captured["tokenizer"] = tokenizer - captured["messages"] = messages - captured["kwargs"] = kwargs + captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs}) return "" monkeypatch.setattr( @@ -248,8 +274,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): ) ) assert out == ["hi"] - # The toggled kwargs must reach the chat-template helper. - assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] - assert captured["kwargs"]["enable_thinking"] is True - assert captured["kwargs"]["reasoning_effort"] == "medium" - assert captured["kwargs"]["preserve_thinking"] is True + # The toggled kwargs must reach the chat-template helper on the real render + # (one of the calls carries the tools; the fallback probe passes tools=None). + tool_renders = [ + c + for c in captured_calls + if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}] + ] + assert tool_renders, captured_calls + render = tool_renders[0] + assert render["kwargs"]["enable_thinking"] is True + assert render["kwargs"]["reasoning_effort"] == "medium" + assert render["kwargs"]["preserve_thinking"] is True diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index 9cf2a62c39..300eb587b3 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -5,8 +5,8 @@ Covers: * GGUF variant listing computes update_available from the already-fetched sibling metadata instead of a second Hub call. - * hf_hub_download_with_xet_fallback(force_download=True) bypasses the - try_to_load_from_cache cache-first early-return. + * hf_hub_download_with_xet_fallback forwards force_download through the shim to the + shared unsloth_zoo helper (which owns the cache-first early-return and its bypass). The cache "Update" action now runs through the download manager as a normal managed download (so it shows in the Downloads panel with progress + cancel), @@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── -def test_force_download_bypasses_cache_first_early_return(monkeypatch): - """force_download=True skips the try_to_load_from_cache early-return and - proceeds to the real download path; force_download=False returns the cached - path without ever attempting a download (X2/F2).""" - import huggingface_hub as hf +def test_force_download_is_forwarded_through_the_shim(monkeypatch): + """The shim's contract is to forward force_download unchanged to the shared helper (which owns the + cache-first early-return and bypass). Verify both False and True reach it (X2/F2).""" import utils.hf_xet_fallback as X - cached_path = "/cache/blob/cached.gguf" + seen = [] - # Pretend the blob IS cached on disk (try_to_load_from_cache is imported - # inside the function from huggingface_hub, and os.path.exists must agree). - monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False) - monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False) + def fake_shared(repo_id, filename, token, **kwargs): + seen.append(kwargs.get("force_download")) + return "/downloaded/path" - attempts = [] + monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) - def fake_attempt(repo_id, filename, token, **kwargs): - attempts.append( - {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")} - ) - return ("ok", "/freshly/downloaded/path") - - monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True) - - # force_download=False: cache-first early-return, no download attempt. - out = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = False ) - assert out == cached_path - assert attempts == [] # never reached the real download - - # force_download=True: bypass the early-return, run the real download. - out2 = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = True ) - assert out2 == "/freshly/downloaded/path" - assert len(attempts) == 1 - assert attempts[0]["force"] is True + assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged # ── multi-revision GGUF blob comparison and update reclaim ── diff --git a/studio/backend/tests/test_native_template_trust_remote_code.py b/studio/backend/tests/test_native_template_trust_remote_code.py new file mode 100644 index 0000000000..60dc80f64c --- /dev/null +++ b/studio/backend/tests/test_native_template_trust_remote_code.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for trust_remote_code in the native-template fallback. + +``render_native_template`` re-fetches a model's native chat template from its +repo when an Unsloth override template (mistral, gemma-4) dropped the tools +schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo +carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use +that same consent or transformers raises (it requires ``trust_remote_code`` to +instantiate a custom tokenizer class), the ``except`` swallows it, and the +request silently keeps the tool-dropping prompt even though the user already +consented to remote code for the model load. + +These tests pin that the stored ``trust_remote_code`` is threaded to the reload, +that the reload is skipped (returns ``None`` without executing code) when no +consent is stored, and that both backend ``model_info`` dicts persist the flag at +load time so the read lands on a value ``load_model`` actually set. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the +# transformers import deferred inside the function). Load it directly so the test +# runs without importing the heavy ``core.inference`` package (unsloth / torch). +_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py" +_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH) +chat_template_helpers = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(chat_template_helpers) + +render_native_template = chat_template_helpers.render_native_template + + +# A native template that emits a tools section only when tools are provided, so the +# with-tools vs no-tools render differs and ``render_native_template`` accepts it. +_NATIVE_TEMPLATE = ( + "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}" + "{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}" + "{% if add_generation_prompt %}assistant:{% endif %}" +) + +_MESSAGES = [{"role": "user", "content": "what is the weather"}] +_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}] + + +class _JinjaTokenizer: + """Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``. + + Stands in for the live model tokenizer that ``render_native_template`` shallow- + copies and re-points at the native template before rendering. + """ + + def __init__(self, chat_template): + self.chat_template = chat_template + + def apply_chat_template( + self, + messages, + tokenize = False, + add_generation_prompt = True, + tools = None, + **kwargs, + ): + from jinja2 import BaseLoader, Environment + env = Environment(loader = BaseLoader()) + return env.from_string(self.chat_template).render( + messages = messages, + tools = tools, + add_generation_prompt = add_generation_prompt, + ) + + +def _install_custom_code_tokenizer(monkeypatch): + """Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise + unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the + native template. Records the ``trust_remote_code`` it was called with.""" + pytest.importorskip("jinja2") + from transformers import AutoTokenizer + + calls = {} + + def fake_from_pretrained( + model_id, + *args, + trust_remote_code = False, + token = None, + **kwargs, + ): + calls["trust_remote_code"] = trust_remote_code + calls["model_id"] = model_id + calls["token"] = token + if not trust_remote_code: + # Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code: + # has_remote_code and not has_local_code and not trust_remote_code -> ValueError. + raise ValueError( + f"The repository {model_id} contains custom code which must be executed " + "to correctly load the model. Please pass the argument " + "`trust_remote_code=True` to allow custom code to be run." + ) + return _JinjaTokenizer(_NATIVE_TEMPLATE) + + monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained)) + return calls + + +def _model_info(trust_remote_code): + return { + "native_chat_template": None, # force the repo reload path + "base_model": None, # non-LoRA: template_source == active_model_name + "trust_remote_code": trust_remote_code, + # Live tokenizer that gets shallow-copied + re-pointed at the native template. + "tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"), + } + + +def test_native_reload_passes_stored_trust_remote_code(monkeypatch): + """With ``trust_remote_code`` stored on ``model_info`` the custom-code reload + succeeds and the tools-advertising native prompt is returned. This FAILS before + the fix (reload omits the flag, raises, is swallowed, returns None).""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = True) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is not None, "native fallback should render the tools prompt with consent" + assert "[AVAILABLE_TOOLS]" in out + assert "get_weather" in out + assert calls["trust_remote_code"] is True # the stored consent was threaded through + # A successful fetch is cached so the next tool turn skips the reload. + assert model_info["native_chat_template"] == _NATIVE_TEMPLATE + + +def test_native_reload_without_consent_returns_none(monkeypatch): + """Without stored consent the custom-code reload raises, is swallowed, and + ``render_native_template`` returns None (no unconsented code execution). Proves + the stored flag -- not a hard-coded True -- drives the reload.""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = False) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is None + assert calls["trust_remote_code"] is False + # A failed fetch must not be cached as "no template" (would pin the tool drop). + assert model_info["native_chat_template"] is None + + +def test_backend_model_info_persists_trust_remote_code(): + """Both backends must store ``trust_remote_code`` on their per-model info dict so + ``render_native_template`` can source the consent value. Guards against the read + landing on a key ``load_model`` never sets (which would silently no-op the fix).""" + inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text() + mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text() + assert '"trust_remote_code": trust_remote_code,' in inf + assert '"trust_remote_code": trust_remote_code,' in mlx diff --git a/studio/backend/tests/test_pr5624_regressions.py b/studio/backend/tests/test_pr5624_regressions.py new file mode 100644 index 0000000000..4f5471675c --- /dev/null +++ b/studio/backend/tests/test_pr5624_regressions.py @@ -0,0 +1,1011 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Regression tests for PR #5624 (DeepSeek R1/V3.x, GLM 4.x, Kimi K2 tool +parsing). Each test pins a specific edge case surfaced during the +review: + +* GLM string-vs-JSON-encoded value coercion (template emits strings + raw and non-strings JSON-encoded; the parser must not coerce a + bare string ``"42"`` into ``42``). +* GLM ```` containing a literal ``<`` (e.g. ``if x < 10``). +* Kimi K2 dotted name ``functions.my.tool:0`` keeps its full name + (``my.tool``) after stripping only the ``functions.`` prefix and + ``:idx`` suffix, while the full id is preserved on the call. +* Kimi K2 bare-counter id (no ``functions.`` prefix, no ``:IDX``) is + dropped rather than surfaced under a numeric name. +* DeepSeek V3.1 truncated mid-stream produces an empty result without + raising. +* ``routes.inference._strip_tool_xml`` strips the DeepSeek envelope and + the Kimi section markers added by this PR. +""" + +import json + +import pytest + +from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_tool_markup, +) + + +# GLM string-vs-JSON-encoded value coercion (finding B in plan) + + +@pytest.mark.parametrize( + "raw_val, expected_python", + [ + # Bare numeric / bool / null shapes are still treated as JSON + # literals (ambiguous with strings; the template doesn't tell us). + ("42", 42), + ("true", True), + ("false", False), + ("null", None), + ("3.14", 3.14), + ("-7", -7), + ("1e3", 1000.0), + ], +) +def test_glm_numeric_and_bool_literals_are_json_decoded(raw_val, expected_python): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == expected_python + + +@pytest.mark.parametrize( + "raw_val", + [ + "hello world", # plain prose + "True", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "None", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "if x < 10: pass", # code with literal < (well, < not in arg_value here) + "{not valid json", # looks like an object but is malformed -- must stay raw + "[oops", # looks like an array but is malformed + ], +) +def test_glm_non_json_shapes_stay_raw(raw_val): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == raw_val + assert isinstance(args["v"], str) + + +def test_glm_json_object_arg_decoded(): + text = ( + "nest\n" + "opts\n" + '{"limit": 10}\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["opts"] == {"limit": 10} + + +def test_glm_json_array_arg_decoded(): + text = ( + "nest\n" + "ids\n" + "[1, 2, 3]\n" + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["ids"] == [1, 2, 3] + + +def test_glm_arg_value_with_literal_less_than(): + text = ( + "run\n" + "code\n" + "if x < 10: pass\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "if x < 10: pass" + + +# GLM 4.7 no-newline emission shape + + +def test_glm_4_7_no_newlines_between_name_and_arg_key(): + """GLM 4.7 strips the ``\\n`` after the name (``{{- ... -}}`` in the + template) so ```` follows directly. Parser must accept both.""" + text = ( + "get_weather" + "cityLondon" + "unitscelsius" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"city": "London", "units": "celsius"} + + +def test_glm_4_7_no_newlines_multi_call(): + """Back-to-back GLM 4.7 calls without intervening newlines.""" + text = ( + "ax1" + "by2" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + + +def test_glm_4_7_does_not_break_qwen_path(): + """Qwen ``{json}`` still dispatches to Qwen; GLM's + first-char ``[^\\n<{]`` excludes ``{``.""" + text = '{"name":"web_search","arguments":{"q":"x"}}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + +# Kimi K2 dotted name + bare counter (finding C in plan) + + +def test_kimi_dotted_namespace_keeps_full_dotted_name(): + # A dotted Kimi id keeps its FULL name; only the ``functions.`` prefix and ``:idx`` suffix drop (vLLM parity). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.my.tool:0" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "my.tool" + assert calls[0]["id"] == "functions.my.tool:0" + + +def test_kimi_two_sections_in_one_stream_both_parse(): + """Outer loop walks every ``<|tool_calls_section_begin|>...end|>`` + so vLLM / SGLang parity holds even on multi-section streams.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"x":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " some prose between sections " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.b:0" + '<|tool_call_argument_begin|>{"y":2}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + assert calls[0]["id"] == "functions.a:0" + assert calls[1]["id"] == "functions.b:0" + + +def test_kimi_bare_counter_id_is_dropped(): + """Bare-digit id (``3``) is dropped (matches vLLM); SGLang infers + name from schema, which we don't have at parse time.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>3" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +# DeepSeek truncated mid-stream + + +def test_deepseek_v3_1_huge_truncated_body_is_linear(): + """Adversarial input: DeepSeek envelope with no JSON brace and a + 50k-char body. A regex-based ``[^\\n<]+?`` name capture is O(N^2) + here; the parser uses ``str.find`` on the sep marker so it stays + linear. Budget 1s to flag any future regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>fn<｜tool▁sep｜>" + "x" * 50_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"V3 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_deepseek_r1_huge_fenceless_body_is_linear(): + """R1 detection used a greedy ``([^\\n]+)\\n```json`` regex that is O(N^2) on a + fence-less body of repeated ``function`` tokens. The parser now scans with + ``str.find``; budget 1s to flag any regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>a" * 40_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"R1 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_glm_unclosed_body_many_arg_keys_is_linear(): + """An unclosed GLM ```` body runs to EOF; a lazy-group ``finditer`` + over many bare ```` tokens was O(N^2). The parser now walks pairs with + ``str.find``; budget 1s.""" + import time as _time + + text = "foo\n" + "k" * 40_000 + start = _time.time() + parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"GLM path is non-linear: {elapsed:.2f}s" + + +def test_deepseek_r1_fenced_json_parses(): + """R1 wraps args in a ```json fence after ``functionNAME``.""" + import json as _json + + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC","unit":"c"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert _json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC", "unit": "c"} + + +def test_deepseek_v3_1_truncated_arguments_drops_call_without_crash(): + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"' # no closing brace, no end markers + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +def test_deepseek_v3_1_truncated_after_end_marker_still_yields_call(): + text = ( + "<｜tool▁calls▁begin｜>" "<｜tool▁call▁begin｜>get_time" "<｜tool▁sep｜>" '{"city":"Tokyo"}' + # neither <｜tool▁call▁end｜> nor <｜tool▁calls▁end｜> + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +# Routes-layer strip across the three new families + + +def test_routes_layer_strip_removes_deepseek_envelope(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_kimi_section(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_glm_block(): + """``.*?`` covers GLM via the Qwen pattern.""" + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "web_search\n" + "q\nx\n" + "" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +# strip_tool_markup (parser-level finalise path) over the new families + + +def test_strip_tool_markup_handles_deepseek_envelope(): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "｜tool▁" not in stripped + assert "get_time" not in stripped and "Tokyo" not in stripped + + +def test_strip_tool_markup_handles_kimi_section(): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "tool_calls_section_begin" not in stripped + + +# Round-2 review findings: GLM quoted-string / unclosed-arg, DeepSeek +# strict terminator, nested wrapper-less Gemma strip + + +def test_glm_quoted_string_arg_keeps_its_quotes(): + # A GLM string value emitted verbatim that itself begins with a quote. + text = ( + "web_search\n" + "query\n" + '"exact phrase"\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == '"exact phrase"' + + +def test_glm_unclosed_arg_value_is_rejected_in_strict_mode(): + # Closing present but a value never closes: strict mode must reject + # the whole call rather than execute it with the argument silently dropped. + text = ( + "web_search\n" + "query\n" + "Tokyo weather" # no + "" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # With Auto-Heal the partial value is kept, not dropped to a no-arg call. + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + args = json.loads(healed[0]["function"]["arguments"]) + assert "Tokyo weather" in args.get("query", "") + + +def test_deepseek_v3_missing_call_terminator_rejected_in_strict_mode(): + # Envelope closes but the per-call <｜tool▁call▁end｜> is absent. Strict mode + # must reject (it is truncated/merged); Auto-Heal still parses it. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁calls▁end｜>" # envelope end only, no per-call end + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + assert healed[0]["function"]["name"] == "get_time" + + +def test_deepseek_v3_with_call_terminator_parses_in_strict_mode(): + # Sanity: a well-formed V3 call (with the per-call end marker) still parses + # under strict mode after the terminator check. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + + +def test_strip_tool_markup_removes_nested_wrapperless_gemma_call(): + # Wrapper-less Gemma call with a NESTED object arg: the balanced helper must strip the whole call, not leave a trailing ``}``. + text = "answer: call:f{loc:{city:NYC},n:3} done" + stripped = strip_tool_markup(text, final = True) + assert "call:f" not in stripped + assert "}" not in stripped + assert "answer:" in stripped and "done" in stripped + + +# Pass-3 review findings: bare-Kimi streaming (non-final) strip symmetry +# and the wrapper-less Gemma route-display strip + + +def test_strip_tool_markup_non_final_removes_bare_kimi_call(): + # A bare ``<|tool_call_begin|>...<|tool_call_end|>`` (no section wrapper): the CLOSED (final=False) strip must remove it too. + text = ( + "before " + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = False) + assert "tool_call_begin" not in stripped + assert "tool_call_end" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_routes_layer_strip_removes_wrapperless_gemma_call(): + # Gemma 4 (skip_special_tokens) emits a wrapper-less ``call:NAME{..}`` with no XML markers. + from routes.inference import _strip_tool_xml as _routes_strip + + text = 'before call:web_search{query:"weather in Sydney"} after' + stripped = _routes_strip(text) + assert "call:web_search" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_deepseek_envelope_end_inside_arg_string_is_not_a_truncation(): + # A DeepSeek V3.1 call whose argument string contains the literal envelope-end token must not be dropped. + content = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>web_search<｜tool▁sep｜>" + '{"query":"what does <｜tool▁calls▁end｜> mean"}' + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "what does <｜tool▁calls▁end｜> mean" + } + + +def test_glm_value_containing_literal_arg_value_close_is_preserved(): + # A GLM string argument may legitimately contain . + content = ( + "runcode" + 'print("")' + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + +def test_attribute_form_function_with_embedded_marker_runs_outer_call(): + # is a supported envelope; a DeepSeek/Kimi marker inside one of its + # parameter values is data, not a second call. + content = ( + '' + "The Kimi format is <|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["respond"], calls + + +def test_wrapperless_gemma_call_gated_by_enabled_tools(): + # Once skip_special_tokens removes the <|tool_call> wrapper, call:NAME{...} is + # indistinguishable from prose documenting the Gemma syntax. + prose = "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + assert parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) == [] + # The display strip is gated the same way, so the example survives in the answer. + assert "call:foo{x:1}" in strip_tool_markup( + prose, final = True, enabled_tool_names = {"web_search"} + ) + # An enabled name is still a real call (parsed, and stripped from display). + real = "Answer. call:web_search{query:hi}" + calls = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert "call:web_search" not in strip_tool_markup( + real, final = True, enabled_tool_names = {"web_search"} + ) + + +def test_kimi_section_end_inside_arg_string_is_not_a_truncation(): + # In a multi-call Kimi section, a later call whose argument holds the literal section-end token must not truncate the section. + content = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>" + '{"q":"cats"}<|tool_call_end|>' + "<|tool_call_begin|>functions.explain:1<|tool_call_argument_begin|>" + '{"text":"the token <|tool_calls_section_end|> means end"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["search", "explain"], calls + assert json.loads(calls[1]["function"]["arguments"]) == { + "text": "the token <|tool_calls_section_end|> means end" + } + + +def test_closed_envelope_before_deepseek_block_owns_turn(): + # Document order is the contract: a CLOSED / call that precedes a + # DeepSeek/Kimi block owns the turn, even when prose frames it as an example. + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>search_web\n" + "```json\n" + '{"query":"weather in Paris"}\n' + "```" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + prose = ( + 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' + ) + calls = parse_tool_calls_from_text(prose + deepseek) + assert [c["function"]["name"] for c in calls] == ["example_tool"], calls + + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.lookup:0" + '<|tool_call_argument_begin|>{"id":7}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls_k = parse_tool_calls_from_text("Example: {} and now:\n" + kimi) + assert [c["function"]["name"] for c in calls_k] == ["demo"], calls_k + + +def test_marker_inside_closed_outer_envelope_still_runs_outer_call(): + # The guard must fire when the marker sits INSIDE a closed outer / envelope's arguments: the OUTER call wins. + outer = ( + "what does <｜tool▁calls▁begin｜> mean" + ) + calls = parse_tool_calls_from_text(outer) + # The outer envelope is the real call; the embedded DeepSeek marker must not + # hijack the parse into a spurious tool. + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "q": "what does <｜tool▁calls▁begin｜> mean" + } + + +def test_truncated_outer_envelope_with_embedded_marker_heals_outer_call(): + # A TRUNCATED outer call embedding a DeepSeek/Kimi marker in its argument still Auto-Heals as the outer call. + trunc = 'x = "<｜tool▁calls▁begin｜>sample"' + calls = parse_tool_calls_from_text(trunc) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_python_tag_call_with_embedded_marker_runs_outer_call(): + # ``<|python_tag|>`` is Llama-3's tool-call envelope, so a DeepSeek/Kimi example quoted + # in its argument is data: the OUTER python_tag call (``web_search``) must run, not the + # embedded marker (``delete_all``). + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>delete_all<｜tool▁sep｜>{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + for embedded in (kimi, deepseek): + builtin = '<|python_tag|>web_search.call(query="explain ' + embedded + '")' + calls = parse_tool_calls_from_text(builtin, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + custom = ( + '<|python_tag|>{"name":"web_search","parameters":' + '{"query":"explain ' + embedded + '"}}' + ) + calls = parse_tool_calls_from_text(custom, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A bare ``<|python_tag|>`` prose mention (no call shape) must NOT be treated as an + # envelope: a real Kimi call after it still parses (the call-shaped lookahead guard). + prose = "The token <|python_tag|> is used. " + kimi + calls = parse_tool_calls_from_text(prose) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +def test_gemma_wrapperless_quoted_value_with_comma_not_split(): + # A wrapper-less Gemma call whose quoted value contains ``, key:``. + text = 'call:web_search{query:"weather, location: Boston", limit:3}' + 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": "weather, location: Boston", + "limit": 3, + } + + +def test_literal_close_tag_in_xml_arg_before_marker_runs_outer_call(): + # A literal ```` inside an outer XML argument (before a marker) is not the envelope close: the span reaches the REAL final close. + text = ( + 'x = " ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_literal_tool_call_close_in_qwen_json_before_marker_runs_outer_call(): + # A Qwen/Hermes whose JSON argument holds a literal then a marker must run the OUTER call. + text = ( + '{"name":"search","arguments":{"query":"explain then ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["search"], calls + # Back-to-back Qwen calls still parse independently (real-close span must keep the + # negative-lookahead that separates adjacent calls). + bb = ( + '{"name":"a","arguments":{}}' + '{"name":"b","arguments":{}}' + ) + assert [c["function"]["name"] for c in parse_tool_calls_from_text(bb)] == ["a", "b"] + + +def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): + # DeepSeek R1 multi-call where the FIRST call has balanced JSON but omits its close + # fence/terminator, followed by a well-formed second call. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n```json\n" + '{"city":"SF"}\n```' # no <｜tool▁call▁end｜> + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time\n```json\n" + '{"tz":"UTC"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + heal = [c["function"]["name"] for c in parse_tool_calls_from_text(text)] + assert "get_time" in heal, heal + # Strict keeps the later well-formed call; heal must be a superset. + strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert set(strict) <= set(heal), (strict, heal) + + +def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): + # A wrapper-less Gemma call whose quoted argument mentions another enabled tool must not execute that nested name. + text = 'call:web_search{query:"explain call:delete_all{target:files}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "explain call:delete_all{target:files}" + } + # Two genuinely separate calls still both parse. + two = "call:web_search{query:hi}call:get_time{tz:UTC}" + assert [ + c["function"]["name"] + for c in parse_tool_calls_from_text(two, enabled_tool_names = {"web_search", "get_time"}) + ] == ["web_search", "get_time"] + + +def test_leading_bare_json_call_owns_quoted_gemma_snippet(): + # Document order: a leading Llama-3.2 bare-JSON call with trailing prose owns the turn. + text = ( + '{"name":"lookup","parameters":{"note":"use call:web_search{query:cats} for this"}}\n' + "That is the call I would make." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "note": "use call:web_search{query:cats} for this" + } + + # Same with the ``;`` inter-call separator: both real calls parse, the + # quoted snippet still does not. + two = ( + '{"name":"lookup","parameters":{"note":"see call:web_search{query:cats}"}};' + '{"name":"lookup","parameters":{"q":"second"}}' + ) + calls_two = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls_two] == ["lookup", "lookup"], calls_two + + +def test_leading_gemma_call_still_wins_over_trailing_json_example(): + # Reverse control: a real leading Gemma call followed by a bare-JSON example keeps the Gemma call (bare JSON matches only a LEADING object). + text = 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "demo_tool"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # And prose-only enabled Gemma syntax (no leading JSON) still promotes: the + # markerless by-design behaviour is unchanged. + prose = "You can run call:web_search{query:cats} to search." + calls_p = parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_p] == ["web_search"], calls_p + + +def test_leading_gemma_call_owns_quoted_mistral_trigger(): + # A leading wrapper-less Gemma call whose argument quotes a Mistral trigger must win: the [TOOL_CALLS] literal is data. + text = 'call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # Reverse control: a real leading Mistral call still parses normally. + real = '[TOOL_CALLS]delete_all{"x":1}' + calls_m = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["delete_all"], calls_m + + # A DISABLED Gemma example quoting the trigger is dropped as prose and a + # real call after it still parses (drop-the-span recursion). + mixed = ( + 'Example: call:demo{note:"see [TOOL_CALLS]delete_all{}"}\n' + '[TOOL_CALLS]web_search{"q":"real"}' + ) + calls_d = parse_tool_calls_from_text(mixed, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_d] == ["web_search"], calls_d + + +def test_chained_bare_json_owns_kimi_marker_in_later_call(): + # Document order: two ;-chained bare-JSON calls own the turn even when the second's argument quotes a complete Kimi snippet. + kimi = ( + "<|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + ) + two = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"note":"' + kimi + '"}}' + ) + calls = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + # Reverse control: prose followed by a real Kimi block still parses. + real = "Let me check.\n<|tool_calls_section_begin|>" + kimi + "<|tool_calls_section_end|>" + calls_k = parse_tool_calls_from_text(real, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_k] == ["delete_all"], calls_k + + # A closed leading Mistral call preceding a trailing Kimi example owns the + # turn too (same closed-call-precedes-marker rule). + mistral = '[TOOL_CALLS]lookup{"q":"first"} then example ' + kimi + calls_m = parse_tool_calls_from_text(mistral, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["lookup"], calls_m + + +def test_nested_gemma_values_keep_commas_and_parens(): + # Nested wrapper-less Gemma mappings/arrays use the top-level delimiter rules, so nested arguments are not split. + calls = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2),lang:py}}", enabled_tool_names = {"python"} + ) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "opts": {"code": "print(1,2)", "lang": "py"} + } + + arr = parse_tool_calls_from_text( + "call:python{opts:[1,2,{a:f(1,2)}]}", enabled_tool_names = {"python"} + ) + assert json.loads(arr[0]["function"]["arguments"]) == {"opts": [1, 2, {"a": "f(1,2)"}]} + + prose_comma = parse_tool_calls_from_text( + "call:python{opts:{note:hello, world}}", enabled_tool_names = {"python"} + ) + assert json.loads(prose_comma[0]["function"]["arguments"]) == {"opts": {"note": "hello, world"}} + + quoted = parse_tool_calls_from_text( + 'call:python{opts:{q:say "a, b" now,n:3}}', enabled_tool_names = {"python"} + ) + assert json.loads(quoted[0]["function"]["arguments"]) == { + "opts": {"q": 'say "a, b" now', "n": 3} + } + + # Controls: nested quoted values and multi-key mappings are unchanged, and + # a truncated nested value still falls back to the raw string. + nested_q = parse_tool_calls_from_text( + 'call:python{loc:{city:"New York"}}', enabled_tool_names = {"python"} + ) + assert json.loads(nested_q[0]["function"]["arguments"]) == {"loc": {"city": "New York"}} + multi = parse_tool_calls_from_text( + "call:python{opts:{a:1,b:2},n:3}", enabled_tool_names = {"python"} + ) + assert json.loads(multi[0]["function"]["arguments"]) == {"opts": {"a": 1, "b": 2}, "n": 3} + trunc = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2}}", enabled_tool_names = {"python"} + ) + assert json.loads(trunc[0]["function"]["arguments"]) == {"opts": "{code:print(1,2}"} + + +def test_multi_gemma_calls_own_turn_over_signal_in_later_call(): + # Document order: when the first enabled Gemma call closes before the first foreign signal, the leading call still owns the turn. + en = {"get_time", "web_search", "delete_all"} + both = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in both] == ["get_time", "web_search"], both + assert json.loads(both[1]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # XML and Kimi markers in the later call's strings stay data too. + xml = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see delete_all"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml] == ["get_time", "web_search"], xml + kimi = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see <|tool_call_begin|>' + 'functions.delete_all:0<|tool_call_argument_begin|>{}<|tool_call_end|>"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in kimi] == ["get_time", "web_search"], kimi + + # A trailing prose example after the closed leading call defers the same way. + prose = parse_tool_calls_from_text( + "call:get_time{} Example: [TOOL_CALLS]delete_all{}", enabled_tool_names = en + ) + assert [c["function"]["name"] for c in prose] == ["get_time"], prose + + +def test_multi_gemma_ownership_reverse_controls(): + # A real leading Mistral/XML call with a trailing Gemma example keeps the leading call; a signal before every Gemma call keeps normal order. + en = {"get_time", "web_search", "delete_all"} + mistral = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"delete_all","arguments":{}}] Example: call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in mistral] == ["delete_all"], mistral + xml_first = parse_tool_calls_from_text( + '{"name":"delete_all","arguments":{}} call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml_first] == ["delete_all"], xml_first + agnostic = parse_tool_calls_from_text( + 'call:foo{} {"name":"delete_all","arguments":{}}' + ) + assert [c["function"]["name"] for c in agnostic] == ["delete_all"], agnostic + + +def test_disabled_leading_bare_json_does_not_hide_later_marker_call(): + # A leading bare-JSON object with a NOT-enabled name is prose: the real DeepSeek/Kimi call after it still parses. + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"cats"}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + kimi, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"q":"cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls_ds = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + deepseek, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls_ds] == ["web_search"], calls_ds + + +def test_disabled_leading_bare_json_ownership_controls(): + kimi_delete = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + # ENABLED leading name still owns the turn (document order, the shipped + # inside-or-after rule). + owns = parse_tool_calls_from_text( + '{"name":"web_search","parameters":{"q":"first"}} ' + kimi_delete, + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert [c["function"]["name"] for c in owns] == ["web_search"], owns + # A marker INSIDE the disabled object's own strings stays data: the span + # is prose, the tail holds no call, so nothing parses. + inside = parse_tool_calls_from_text( + '{"name":"draft","parameters":{"note":"see <|tool_call_begin|>functions.delete_all:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}\nsome trailing prose', + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert inside == [], inside + # Nameless leading JSON answers keep recursing to the real call. + nameless = parse_tool_calls_from_text( + '{"answer":42} ' + kimi_delete, enabled_tool_names = {"delete_all"} + ) + assert [c["function"]["name"] for c in nameless] == ["delete_all"], nameless + # Name-agnostic path unchanged: the leading object is the call. + agnostic = parse_tool_calls_from_text('{"name":"draft","parameters":{}} ' + kimi_delete) + assert [c["function"]["name"] for c in agnostic] == ["draft"], agnostic + + +def test_leading_json_answer_with_prose_keeps_quoted_gemma_snippet_as_data(): + # A LEADING JSON answer followed by prose is data (same contract as the whole-content JSON exemption). + obj = '{"summary":"use call:web_search{query:cats} to search"}\nHope that helps!' + assert parse_tool_calls_from_text(obj, enabled_tool_names = {"web_search"}) == [] + arr = '["use call:web_search{query:cats} to search"]\nHope that helps!' + assert parse_tool_calls_from_text(arr, enabled_tool_names = {"web_search"}) == [] + assert strip_tool_markup(obj, enabled_tool_names = {"web_search"}) == obj + + # A REAL call in the tail after the answer still parses (and strips). + tail = '{"summary":"done"}\ncall:web_search{query:cats}' + calls = parse_tool_calls_from_text(tail, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A leading brace run that is NOT valid JSON gets no exemption. + not_json = "{not json} call:web_search{query:cats}" + calls_nj = parse_tool_calls_from_text(not_json, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_nj] == ["web_search"], calls_nj + + +def test_glm_heal_bounds_unclosed_value_at_tool_call_close(): + # Auto-Heal: a value missing only its before the block's heals to the + # value text, not the close tag and everything after it swallowed into the argument. + one = "get_weathercityNYC" + calls = parse_tool_calls_from_text(one, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["get_weather"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC"} + + # Trailing prose after the close stays out of the healed value. + two = one + "\nLet me check that for you." + calls_two = parse_tool_calls_from_text(two, allow_incomplete = True) + assert json.loads(calls_two[0]["function"]["arguments"]) == {"city": "NYC"} + + # Strict mode still rejects the unclosed value outright. + assert parse_tool_calls_from_text(one, allow_incomplete = False) == [] + + # A value truncated at EOF (no structural tag follows) keeps the partial heal, and a proper + # close whose value holds a literal is untouched by the bounding. + eof = "get_weathercityNew York Ci" + calls_eof = parse_tool_calls_from_text(eof, allow_incomplete = True) + assert json.loads(calls_eof[0]["function"]["arguments"]) == {"city": "New York Ci"} + lit = ( + "get_weathercity" + 'print("")' + ) + calls_lit = parse_tool_calls_from_text(lit, allow_incomplete = True) + assert json.loads(calls_lit[0]["function"]["arguments"]) == {"city": 'print("")'} + + +def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): + # False-alarm literals: the trailing strip arms require a call-shaped + # lookahead, so an answer documenting a marker keeps its tail. + from core.inference.tool_call_parser import strip_tool_markup + + for text in [ + "The Kimi marker <|tool_calls_section_begin|> starts a section.", + "DeepSeek uses <｜tool▁calls▁begin｜> to open calls.", + "See <|tool_call_begin|> in the docs.", + ]: + assert strip_tool_markup(text, final = True) == text + + # Truncated REAL calls still drop, and a bare marker at EOF is a fragment. + truncated_kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q' + ) + assert strip_tool_markup(truncated_kimi, final = True) == "" + assert strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index a7ceb49ed9..89a0b3879b 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,126 @@ class TestTranslatedMessagesValidate: ChatMessage(**m.model_dump(exclude_none = True)) +# reasoning_prefilled mode: Qwen3/GLM enable_thinking templates prefill an unclosed , so +# generation begins inside the think block and emits only the closing ; the 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 into the drawer (covers enable_thinking_effort full-tag output). + 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: GGUF-parity guard -- WITHOUT prefilled, a lone keeps the + # pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and + # every existing caller are byte-identical. + 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) still passes text + # straight through even if reasoning_prefilled were mistakenly set False. + 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..3701a00dd2 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -11,6 +11,8 @@ from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) @@ -127,9 +129,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 now parser-supported, so supports_tools=True +# must hold for all of them; only templates matching none of the five known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -161,27 +162,186 @@ 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 + + +# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks. +# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581). +DEEPSEEK_TEMPLATE = """ +{%- if tools %} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + {%- for tc in message.tool_calls %} + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name + + '<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }} + {%- endfor %} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_deepseek_template_keeps_tools_on(): + """DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE) + assert flags["supports_tools"] is True + + +# GLM 4.5 / 4.6 / 4.7 emit ``NAME\n...... +GLM_TEMPLATE = """ +{%- if tools %} + For each function call, output the function name and arguments within + the following XML format: + {function-name} + {arg-key} + {arg-value} + + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_glm_template_keeps_tools_on(): + """GLM 4.x emits ``NAME\\n...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6") + flags = _detect_safetensors_features(backend, GLM_TEMPLATE) + assert flags["supports_tools"] is True + + +# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks +# with ``functions.NAME:IDX`` as the per-call id. +KIMI_TEMPLATE = """ +{%- if tools %} + <|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|> +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + <|tool_calls_section_begin|> + {%- for tc in message.tool_calls %} + <|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|> + {%- endfor %} + <|tool_calls_section_end|> + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_kimi_template_keeps_tools_on(): + """Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct") + flags = _detect_safetensors_features(backend, KIMI_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 +614,161 @@ 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 + + +@pytest.mark.parametrize( + "opener", + [ + "<｜tool▁calls▁begin｜>", # canonical + "<｜tool_calls_begin｜>", # ASCII underscores + "<｜tool▁calls｜>", # short form + "<｜tool calls begin｜>", # spaces + "<｜tool\\_calls\\_begin｜>", # escaped underscores + ], +) +def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener): + # Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives + # its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ... + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}tools{%- endif %}" + + opener + + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON +# call form, but 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(): + # A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser + # accepts that whitespace via raw_decode. + 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 at + # all -> the pill is still dropped (the gate is not now matching everything). + 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(): + # A template documenting the parser-supported {"function":...} bare-JSON alias + # must keep supports_tools, mirroring the {"name":...} form. + 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 so safetensors/MLX reach +# GGUF reasoning-block parity for enable_thinking models. +class TestSafetensorsReasoningPrefillGate: + # A minimal 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 + # (would otherwise swallow the whole answer 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..4a5423fa87 --- /dev/null +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -0,0 +1,186 @@ +# 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] = [] # sequence of ("reasoning"|"visible"|"tool_start") events + + 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-style enable_thinking_effort: a request with reasoning_effort="none" (and + # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so + # prefilled mode must be OFF. Otherwise the model emits no and a plain + # answer is swallowed whole into reasoning_content, leaving the visible response + # empty (the exact bug: prefilled=True below eats the whole answer). + 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 the corrected prefilled=False, a plain no- answer is + # emitted as visible content rather than swallowed into the thinking drawer. + 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..38b30fe8f6 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-argument indentation survives (str.strip() destroyed it). + 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,15 @@ 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 ```` like the array shape, so the EOS marker is not left + # behind as visible content. + 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 +237,479 @@ class TestParser: == "before " ) + def test_streaming_strip_handles_nested_mistral_json(self): + # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the + # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... + raw = 'ok [TOOL_CALLS]foo{"a":{"b":1}} tail' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_handles_nested_wrapperless_gemma(self): + # Same class of bug for the wrapper-less Gemma call:NAME{...} form with a + # nested object argument. + raw = "ok call:f{loc:{city:NYC},n:3} tail" + out = strip_tool_markup_streaming(raw) + assert "call:f" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): + # A literal ```` in a value is data: the strip must close at the REAL + # ```` and keep trailing prose (the open-ended regex ate to EOF). + 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 emits reasoning as a leading ``[THINK]...[/THINK]`` bracket block + # (not the ```` the reasoning channel renders). The streaming display + # strip must drop it so the raw chain-of-thought does not leak into the + # safetensors content; GGUF routes it to reasoning_content natively. + 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 from the marker on (nothing leaks, and + # the cleaned text only grows as the answer streams in 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 and is 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) == "" + + def test_llama3_python_tag_json_form_non_scalar_args_skipped(self): + # Should NOT fabricate ``{"value": args}`` when the JSON form + # has a non-dict / non-string ``arguments`` value. + for bad in ( + '<|python_tag|>{"name":"foo","arguments":42}', + '<|python_tag|>{"name":"foo","arguments":[1,2,3]}', + '<|python_tag|>{"name":"foo","arguments":null}', + '<|python_tag|>{"name":"foo","arguments":true}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # ── 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 per its training template. + 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 (parser is + # strict about content STARTING 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. Prose like + # ``{"name":"foo","parameters":"a sentence"}`` 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 the guard. + 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, matching the JSON/XML paths and SGLang's base detector. + 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: ``[TOOL_CALLS]name[CALL_ID][ARGS]{json}``. + # The ``[CALL_ID]`` segment must be skipped, not treated as a stop + # (llama.cpp test-chat.cpp:4785 parses this to one call). + 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): + # Magistral wraps reasoning in ``[THINK]...[/THINK]``. A ``[TOOL_CALLS]`` + # inside the reasoning is chain-of-thought, not a real call; 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 merely mentions a tool call but does not emit one + # 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 (after the call) + # 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) == "" + + # ── Gemma 4 wrapper-less (skip_special_tokens stripped) ─────────── + + def test_gemma4_bare_stripped_call(self): + # skip_special_tokens removes <|tool_call>/ and <|"|>, + # leaving a bare call:NAME{...} with an unquoted value. + import json + + text = "call:web_search{query:weather in San Francisco right now}" + 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": "weather in San Francisco right now"} + + def test_gemma4_bare_code_with_commas(self): + # A code value with commas must not truncate at the first comma. + import json + + text = ( + "call:python{code:def f(n):\n a, b = 0, 1\n" + " for _ in range(2, n+1):\n a, b = b, a + b\n" + " return b\n\nprint(f(30))}" + ) + result = parse_tool_calls_from_text(text) + assert result[0]["function"]["name"] == "python" + code = json.loads(result[0]["function"]["arguments"])["code"] + assert "a, b = 0, 1" in code and "print(f(30))" in code + + def test_gemma4_bare_quotes_normalized(self): + # The same value quoted vs unquoted must parse identically so the + # agentic loop can collapse a looping model's repeated calls. + import json + + a = parse_tool_calls_from_text('call:web_search{query:"foo bar"}') + b = parse_tool_calls_from_text("call:web_search{query:foo bar}") + assert json.loads(a[0]["function"]["arguments"]) == {"query": "foo bar"} + assert json.loads(a[0]["function"]["arguments"]) == json.loads( + b[0]["function"]["arguments"] + ) + + def test_gemma4_bare_multi_arg(self): + import json + + text = "call:web_search{query:pytorch latest, url:https://pytorch.org}" + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "pytorch latest", "url": "https://pytorch.org"} + + def test_gemma4_bare_not_matched_in_prose(self): + # A word ending in "call:" must not trigger a bare tool call. + text = "I will recall:that the function{ } is helpful." + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_gemma4_bare_strip_markup_final(self): + text = "Here you go: call:web_search{query:weather today}" + assert "call:web_search" not in 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 @@ -312,6 +808,553 @@ def _make_loop( ), exec_fn +class TestParserDeepSeek: + """DeepSeek R1 / V3 / V3.1 coverage. Markers use full-width pipes + (U+FF5C) and lower-one-eighth-block (U+2581). R1 wraps args in a + Markdown ``` ```json ``` ``` fence; V3 / V3.1 emit bare JSON.""" + + def test_r1_simple_call_with_code_fence(self): + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function" + "<｜tool▁sep｜>special_function\n" + "```json\n" + '{"arg1": 1}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + 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_r1_short_form_outer_marker(self): + # llama.cpp accepts ``<｜tool▁calls｜>`` as the short-form opener. + import json as _json + + text = ( + "<｜tool▁calls｜>function" + "<｜tool▁sep｜>get_time\n" + "```json\n" + '{"city": "Paris"}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_bare_json(self): + # V3 / V3.1 omit the ``function`` prefix and the code fence. + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_v3_1_multi_call_shares_envelope(self): + # Parallel calls share one outer envelope; each inner call has + # its own ``<｜tool▁call▁begin｜>...<｜tool▁call▁end｜>``. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_time" + assert result[1]["function"]["name"] == "get_weather" + + def test_v3_1_with_reasoning(self): + # Reasoning ... precedes the tool block. + text = ( + "I'm thinking\n" + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_strict_rejects_unclosed_envelope(self): + # Envelope truncated mid-stream (no <｜tool▁calls▁end｜>): healed by + # default, rejected with Auto-Heal off. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + ) + assert len(parse_tool_calls_from_text(text)) == 1 + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_v3_1_multi_call_recovers_when_first_end_marker_missing(self): + # First inner call omits its <｜tool▁call▁end｜>; the second must still be parsed. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["get_time", "get_weather"] + + def test_v3_1_strict_recovers_after_missing_call_end(self): + # Strict mode (Auto-Heal off): the FIRST inner call is missing its <｜tool▁call▁end｜> + # terminator, so it is skipped -- but the parser must keep scanning and still return the ... + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "SF"}' + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"tz": "PST"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + # Auto-Heal keeps both; strict skips the truncated first, keeps the second. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == [ + "get_weather", + "get_time", + ] + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_r1_strict_recovers_after_missing_close_fence(self): + # R1 form. + text = ( + "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>get_weather\n```json\n" + '{"city": "SF"}' + "function<｜tool▁sep｜>get_time\n```json\n" + '{"tz": "PST"}' + "\n```<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_deepseek_strip_markup(self): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_deepseek_signal_wakes_streaming(self): + # The streaming buffer state machine must wake on the DeepSeek opener so the rest of the + # section is drained instead of leaked. + text = "<｜tool▁calls▁begin｜>..." + assert has_tool_signal(text) + + def test_deepseek_short_opener_is_stripped(self): + # The short ``<｜tool▁calls｜>`` opener is parsed, so its markup must also be stripped (the + # strip patterns used to require ...calls_begin and left the short-opener markup leaking to ... + text = ( + "before " + "<｜tool▁calls｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + +class TestParserGLM: + """GLM 4.5 / 4.6 / 4.7 coverage. Marker collides with Qwen's + ```` but the body shape is XML kv pairs instead of JSON, + so the dispatch order keeps both formats working.""" + + def test_glm_simple_call(self): + import json as _json + + text = ( + "web_search\n" + "query\n" + "weather Tokyo\n" + "" + ) + 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"]) + # Strings come through raw; the parser does not double-quote. + assert args == {"query": "weather Tokyo"} + + def test_glm_mixed_types_decode_correctly(self): + # Per the chat_template.jinja, strings are emitted raw and non-strings are JSON-encoded. + import json as _json + + text = ( + "complex_function\n" + "name\nJohn Doe\n" + "age\n30\n" + "active\ntrue\n" + "score\n95.5\n" + "" + ) + result = parse_tool_calls_from_text(text) + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"name": "John Doe", "age": 30, "active": True, "score": 95.5} + + def test_glm_multi_call_back_to_back(self): + # GLM emits parallel calls as consecutive ``... + # `` blocks with no outer envelope. + text = ( + "a\nx\n1\n" + "b\ny\n2\n" + ) + 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_glm_unclosed_tool_call_does_not_lose_value(self): + # Truncated mid-stream (no ) -- the parser must + # still surface what it found rather than dropping the call. + text = "web_search\nquery\npartial" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_does_not_break_qwen_path(self): + # Real Qwen emission must still be parsed by the Qwen branch, + # not silently misrouted to GLM (the marker is shared). + text = '{"name":"web_search","arguments":{"q":"x"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_strip_markup(self): + text = ( + "before " + "a\nx\n1\n" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_glm_zero_arg_inline_call(self): + # GLM 4.7 emits a no-argument call inline as ``name`` (name followed + # straight by the close tag, no \n / ). + import json as _json + + text = "get_current_date" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_current_date" + assert _json.loads(result[0]["function"]["arguments"]) == {} + + def test_glm_zero_arg_call_in_parallel_batch(self): + # A no-arg call alongside a normal one must not make either vanish. + text = ( + "get_current_date" + "get_weather\ncity\n" + "Tokyo" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_current_date" + assert result[1]["function"]["name"] == "get_weather" + + def test_glm_string_value_whitespace_preserved(self): + # The template emits string args verbatim, so significant leading / trailing whitespace + # (code, diffs) must survive. + import json as _json + + text = ( + "run\ncode\n" + " indented code " + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"code": " indented code "} + + +class TestParserKimi: + """Kimi K2 / Moonshot coverage. ASCII pipes only (NOT full-width). + Name arrives as ``functions.NAME:IDX``; the parser strips the + prefix and the index to recover the bare callable name while + preserving the full id for round-trip rendering.""" + + def test_kimi_simple_call(self): + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.special_function:0" + "<|tool_call_argument_begin|>" + '{"arg1": 1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + # Bare name recovered; full id preserved verbatim. + assert result[0]["function"]["name"] == "special_function" + assert result[0]["id"] == "functions.special_function:0" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_outer_tool_call_with_embedded_kimi_marker_parses_outer(self): + # A Qwen/Hermes whose argument contains literal Kimi markup (a user asking + # about that syntax) must execute the OUTER call, not the embedded marker via the ... + text = ( + '{"name":"web_search","arguments":{"query":' + '"explain <|tool_call_begin|>functions.evil:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_genuine_kimi_call_without_envelope_still_parses(self): + # Control: a real Kimi call with no leading envelope must + # still go through the pre-pass. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query":"x"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_kimi_multi_call_with_index(self): + # Multiple consecutive calls inside a single section, each + # with its own monotonically incrementing ``:IDX``. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_end|>" + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "read_file" + assert result[0]["id"].endswith(":0") + assert result[1]["function"]["name"] == "web_search" + assert result[1]["id"].endswith(":1") + + def test_kimi_dotted_name_keeps_full_dotted_name(self): + # A dotted Kimi id keeps its FULL name after stripping only the ``functions.`` prefix and + # ``:idx`` suffix -- matching current vLLM ... + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>a.b.c:2" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "a.b.c" + + def test_kimi_dotted_mcp_name_with_functions_prefix(self): + # ``functions.mcp.server-list:0`` must resolve to ``mcp.server-list`` + # (only the ``functions.`` prefix and ``:idx`` are removed). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.mcp.server-list:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp.server-list" + + def test_kimi_multi_call_recovers_when_first_end_marker_missing(self): + # First call omits its <|tool_call_end|>; the second must still parse. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["read_file", "web_search"] + + def test_kimi_handles_unclosed_section(self): + # End marker missing -- the parser must still extract the call. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.foo:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + + def test_kimi_strip_markup(self): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.x:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_kimi_signal_wakes_streaming(self): + text = "<|tool_calls_section_begin|>..." + assert has_tool_signal(text) + + def test_kimi_call_without_section_wrapper(self): + # llama.cpp makes the ``<|tool_calls_section_begin|>`` wrapper optional -- Kimi K2 can emit + # a bare ``<|tool_call_begin|>`` call. + import json as _json + + text = ( + "<|tool_call_begin|>functions.execute_command:0" + "<|tool_call_argument_begin|>" + '{"cmd":"ls"}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "execute_command" + assert _json.loads(result[0]["function"]["arguments"]) == {"cmd": "ls"} + + def test_kimi_malformed_json_recovers_later_calls(self): + # A call with malformed / truncated JSON must not drop the valid calls that follow it in + # the same section (the bad call is skipped, the good one is recovered). + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"city":"Beijing"' # missing closing brace + "<|tool_call_end|>" + "<|tool_call_begin|>functions.b:1" + '<|tool_call_argument_begin|>{"city":"Shanghai"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "b" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Shanghai"} + + +class TestParserCrossFormatRouting: + """Ensure the per-format dispatch order doesn't misroute any + family. Real emissions for each new family + every old family + must still parse correctly when intermixed.""" + + def test_dispatch_routes_each_family_correctly(self): + cases = [ + ( + "Qwen", + '{"name":"a","arguments":{"x":1}}', + "a", + ), + ( + "DeepSeek V3.1", + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>", + "get_time", + ), + ( + "GLM", + "web_search\n" + "q\nx\n" + "", + "web_search", + ), + ( + "Kimi", + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.add:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>", + "add", + ), + ] + for label, text, expected_name in cases: + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"{label}: parser missed the call" + assert result[0]["function"]["name"] == expected_name, ( + f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}" + ) + + def test_all_new_markers_in_tool_xml_signals(self): + # The safetensors / MLX streaming buffer must wake on every supported emission marker -- + # otherwise the BUFFERING state leaks tool content to the user before parse. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_active_tools_are_passed_to_single_turn_after_render_html_success(): captured_tool_names: list[list[str]] = [] exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) @@ -347,6 +1390,198 @@ 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 call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; + # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. + 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 ``custom_tools`` bare form ``{"name":..,"parameters":..}`` carries no + # XML signal. The loop must BUFFER it until the object closes and execute it via + # the safety net, never leaking the raw JSON to streaming clients as content. + 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 (e.g. a person record + # ``{"name":"Alice",...}``) must be shown as the answer, not misread as a call + # to a disabled tool and dropped. _make_loop enables web_search/python/terminal. + 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, so the buffer is + # held open across chunks before 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_gemma_wrapperless_call_is_not_streamed_as_content(): + # Gemma 4 wrapper-less ``call:NAME{...}`` has no XML signal; the loop must hold + # it (BUFFERING) and execute it, never streaming the raw call text. + loop, exec_fn = _make_loop( + turns = [["call:web_search{query:cats}"], ["Found."]], + 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("call:web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_with_whitespace_is_suppressed_when_streamed(): + # Gemma may emit ``call : NAME{...}`` with whitespace around the colon, split across stream + # chunks. + loop, exec_fn = _make_loop( + turns = [["call", " : ", "web_search", "{query:cats}"], ["Found."]], + 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("call" in t for t in contents), contents + + +def test_long_gemma_tool_name_is_not_streamed_as_content(): + # A tool name longer than the small buffer cap (OpenAI 64 chars, MCP longer) + # must still be held: the ``call:NAME`` prefix keeps buffering until ``{`` + # instead of leaking ``call:longname`` as visible text. + long_name = "mcp__github__list_repository_issues" # 35 chars + turns = iter([list('call:%s{repo:"octo/hello"}' % long_name), ["Done."]]) + + def _gen(_messages): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["RESULT"]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": long_name}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [(long_name, {"repo": "octo/hello"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:" in t for t in contents), contents + + +def test_leading_json_answer_is_not_dropped(): + # A leading ``{...}`` that is NOT a tool call must still surface as content: + # the bare-JSON hold can only ever delay it to end-of-object, never drop 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 plan-without-action nudge must name the tools actually enabled, never the + # old hardcoded ``web_search``/``python`` (which a restricted set would reject). + 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 safetensors nudge must stay silent for backend parity + # with the GGUF loop, so only the single 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 +1641,154 @@ 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 agentic loop must recognise Llama-3's <|python_tag|> + # marker, drain the rest of 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 emit a bare-JSON tool call + # ``{"name":..,"parameters":..}`` with NO XML signal. The loop's + # safety-net parse must still fire the tool instead of treating the + # turn as "planned without calling tools" and re-prompting the model + # into giving up. Regression for the has_tool_signal gate that + # dropped these; GGUF's llama-server parses them natively. + 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_deepseek_v3_1_form(self): + # DeepSeek V3.1 emission inside the agentic loop -- the buffer state machine must wake on + # ``<｜tool▁calls▁begin｜>`` and the parser must extract the V3.1 bare-JSON body. + loop, exec_fn = _make_loop( + turns = [ + [ + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>web_search", + "<｜tool▁sep｜>", + '{"query":"Tokyo weather"}', + "<｜tool▁call▁end｜>", + "<｜tool▁calls▁end｜>", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + def test_glm_form(self): + # GLM 4.x emission: ``NAME\n...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "web_search\n", + "query\n", + "Tokyo\n", + "", + ], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + + def test_kimi_form(self): + # Kimi K2 emission ``<|tool_calls_section_begin|>...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>functions.web_search:0", + "<|tool_call_argument_begin|>", + '{"query":"Tokyo"}', + "<|tool_call_end|>", + "<|tool_calls_section_end|>", + ], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + # The bare name must reach execute_tool, even though the model + # emitted ``functions.web_search:0`` as the formatted id. + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + # tool_start carries the original full id so the conversation + # roundtrip can replay it verbatim. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "functions.web_search:0" + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -765,6 +2148,59 @@ 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/disabled no-op turn is a correction turn and must NOT spend the + # caller's tool budget, so with max_tool_iterations=2 the model can still make a + # DISTINCT valid call after repeating one. Only turns that actually execute a + # tool count -- matching the GGUF loop. (The budget used to be charged per + # non-re-prompt iteration, so the duplicate burned the second slot and the third + # turn was sent with no tools, dropping the ``python`` call.) + 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 +2389,253 @@ 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 rather than burn forever. + 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 (< 32 chars + no + # marker prefix). The unified buffer-end path must still + # trigger the intent re-prompt, not silently terminate. + 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: one re-prompt, then one real tool call, + # then the budget-exhausted final answer must still fire. If the + # re-prompt ate the slot the tool call would never run. + 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": "print(1)"}, not + # {"query": ...}, so the python sandbox actually executes 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 state machine must wake on every emission + # marker the shared parser knows -- otherwise Llama-3 / Mistral + # / Gemma 4 emissions slip past as plain prose when the + # llama-server structured channel fails. + 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 function must delegate to the shared + # strip_tool_markup so closed-pair markup is removed for every + # emission family (Llama-3 <|python_tag|>, Mistral [TOOL_CALLS], + # Gemma 4 <|tool_call>...). + 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 ``arguments`` to the same + # per-tool canonical key -- ``code`` for python, ``command`` for + # terminal, ``query`` for everything else. The mapping is centralised in + # the shared ToolLoopController (both backends route bare-string args + # through ``coerce_tool_arguments``), so the two paths cannot drift. + 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 forward-looking + # phrases on both backends so behaviour is the same on Mac (MLX + # / safetensors) and on Linux (GGUF). + 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 + # force a tool-call 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() @@ -1369,6 +3052,28 @@ class TestGuardrails: and event.get("type") in {"tool_start", "tool_end"} ] + def test_same_turn_distinct_calls_are_capped(self): + # >_MAX_TOOL_CALLS_PER_TURN DISTINCT calls in one turn must be capped so a runaway turn + # cannot fan out into many executions (the GGUF path is held back by llama-server's lazy ... + from core.inference.safetensors_agentic import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + turn = "".join( + '{"name":"web_search","arguments":{"query":"q%d"}}' % i + for i in range(n) + ) + loop, exec_fn = _make_loop( + turns = [[turn], ["final"]], + exec_results = ["r"] * n, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert len(exec_fn.calls) == _MAX_TOOL_CALLS_PER_TURN + # The first N distinct queries executed, in document order. + assert [a["query"] for _name, a in exec_fn.calls] == [ + "q%d" % i for i in range(_MAX_TOOL_CALLS_PER_TURN) + ] + def test_coerce_string_args_python_uses_code_key(self): assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} @@ -1407,5 +3112,623 @@ 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 does + # not blow up the entire test file at collection time. + 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: 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 any + # 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 + # should 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 fully 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 around a Llama-3.2 bare-JSON object that uses + # ``parameters`` instead of ``arguments``. The bare-JSON and + # python_tag paths already accept both keys; this path now does + # too. Was extracting name only and silently dropping the args. + 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 ```` attribute form must flip + # the streaming buffer; otherwise the end-of-turn safety-net parse in + # the agentic loop is gated off and the real 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 + # the legacy ```` form. + text = 'result X' + assert strip_tool_markup(text, final = True) == "result" + + def test_llama3_chat_template_round_trip(self): + # Meta's official Llama-3.x chat template prefixes every + # assistant turn with + # ``<|start_header_id|>assistant<|end_header_id|>\n\n``. The + # sentinel-strip in ``_parse_llama3_bare_json`` must reach past + # the role label to the JSON body, else every round-tripped + # tool call in history silently drops. + 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 assistant turn closes with ``<|eot_id|>``, then the + # new header opens. Both sentinels + the 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): + # Models routinely follow a tool call with explanatory prose. + # Body must terminate at ```` even without a + # ```` wrapper, else trailing prose leaks into the + # last parameter 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_render_with_native_template_returns_render_only_when_tools_emitted(): + # The native-template fallback re-renders with the model's repo template when an override drops + # the tools schema. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "") + + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) # never reflects tools + + out = render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert out == "hi|TOOLS=web_search" + # The native template must be restored on the live tokenizer after probing. + assert model_info["tokenizer"].chat_template == "OVERRIDE" + + assert ( + render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = ignoring, + ) + is None + ) + + # No tokenizer and no processor -> return None instead of an AttributeError. + no_tok = {"native_chat_template": "TPL"} + assert ( + render_native_template( + model_info = no_tok, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + is None + ) + + +def test_render_with_native_template_does_not_mutate_shared_tokenizer(): + # The shared tokenizer must never carry the temporary native template, even mid-render: this + # runs outside the generation lock, so a concurrent request could otherwise render with the ... + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + shared = SimpleNamespace(chat_template = "OVERRIDE") + seen = [] + + def capture(tokenizer, msgs, *, tools, **_kw): + seen.append((tokenizer is shared, shared.chat_template)) + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = {"native_chat_template": "TPL", "tokenizer": shared} + render_native_template( + model_info = model_info, + active_model_name = "x", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = capture, + ) + # Rendering happened on a copy, and the shared tokenizer stayed "OVERRIDE" + # throughout (never the temporary "TPL"). + assert seen and all(not is_shared for is_shared, _ in seen) + assert all(tpl == "OVERRIDE" for _, tpl in seen) + assert shared.chat_template == "OVERRIDE" + + +def test_native_template_loads_from_base_model_for_lora(monkeypatch): + # For a LoRA adapter the chat template lives on the base model; active_model_name + # is the adapter id and may ship no template. The loader must read base_model. + from types import SimpleNamespace + + import transformers + + from core.inference.chat_template_helpers import render_native_template + + captured = {} + + def fake_from_pretrained(name, *args, **kwargs): + captured["source"] = name + return SimpleNamespace(chat_template = "BASE_TPL") + + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained) + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = { + "base_model": "base/model-id", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + out = render_native_template( + model_info = model_info, + active_model_name = "adapter/path", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = emitting, + ) + assert captured["source"] == "base/model-id" + assert out == "hi|T" + + +def test_render_with_native_template_fallback_swaps_when_override_drops_tools(): + # The shared gate (used by the transformers and MLX backends): when the live render is + # identical with and without tools, re-render with the native template and return it. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + # apply_fn that IGNORES tools -> live render drops the schema. + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) + + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + # Native render emits the tools, so the fallback swaps to it. + def native_emits(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS" if tools else "") + + out = render_with_native_template_fallback( + formatted_prompt = ignoring(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = lambda tok, msgs, *, tools, **kw: ( + native_emits(tok, msgs, tools = tools) + if getattr(tok, "chat_template", None) == "TPL" + else ignoring(tok, msgs, tools = tools) + ), + ) + assert out == "hi|TOOLS", out + + +def test_render_with_native_template_fallback_keeps_prompt_when_tools_emitted(): + # Live render already differs with vs without tools -> no fallback, returned + # unchanged. Also a no-tools call is a passthrough. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + kept = render_with_native_template_fallback( + formatted_prompt = emitting(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert kept == "hi|T", kept + + # No tools -> passthrough (native template never consulted). + passthrough = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(), + model_info = {}, + active_model_name = "x", + messages = messages, + tools = None, + apply_fn = emitting, + ) + assert passthrough == "hi" + + +def test_render_with_native_template_fallback_keeps_prompt_when_no_tools_probe_raises(): + # A template that REQUIRES tools can raise on the no-tools probe. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def raises_without_tools(tokenizer, msgs, *, tools, **_kw): + if not tools: + raise RuntimeError("template requires tools") + return "".join(m["content"] for m in msgs) + "|T" + + out = render_with_native_template_fallback( + formatted_prompt = "hi|T", + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = raises_without_tools, + ) + assert out == "hi|T", out + + +def test_truncated_bare_json_at_eof_is_not_leaked(): + # Stream ends mid bare-JSON object: the held fragment must be dropped at the + # EOF resolver, not flushed as plain assistant content (GGUF parity). + 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 whose arguments exceed _MAX_BARE_JSON_BUFFER must DRAIN + # (suppress) rather than stream the raw JSON prefix, and still execute once + # the full object is parsed by the safety net. + 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 tool call and must still + # stream -- the oversized DRAIN route is gated on a "name" key. + 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: + # the oversized DRAIN branch was gated only on the presence of a "name" key, so a + # large ordinary record ({"name":"Alice",...}) was drained instead of shown. + 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 ordinary JSON answer whose name is not an enabled tool, held to EOF, + # must be shown -- the EOF bare-JSON DRAIN branch was gated only on a "name" key. + 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 ordinary JSON answer with a NESTED ``"name"`` matching an enabled + # tool ({"result":{"name":"web_search",...) must be shown, not suppressed: the + # gate now extracts the TOP-LEVEL name only, so the nested field is just data. + 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 complete bare-JSON call executes, the assistant content fed to the + # next turn must not contain the raw call (next-turn contamination). + 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(): + # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did + # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is + # preserved), matching the XML strip in the same drain branch. With Auto-Heal ON + # the same fragment is 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 safetensors buffering gate must recognise the "function" bare-JSON alias + # the parser accepts, so a truncated/complete {"function":} call is + # buffered/healed instead of streaming as visible content. + 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 that starts with a literal marker is a false alarm: the + # drain finds no calls and 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 kept content (next-turn assistant + # 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_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py new file mode 100644 index 0000000000..5c298a7966 --- /dev/null +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path. + +The parser and the cumulative-text state machine are already covered exhaustively by +``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the +*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and +``MLXInferenceBackend`` render the prompt through the shared +``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the +shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and +``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU +test of that seam covers the macOS MLX path too). + +This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the +``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full +agentic chain end to end: + + tools injected into the template -> loop parses the call -> tool dispatched once -> + tool result fed back -> generation re-entered -> final answer streamed. + +It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling +end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the +tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704). +""" + +from core.inference.chat_template_helpers import apply_chat_template_for_generation +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +TOOL_NAME = "get_weather" +TOOL_ARGS = {"city": "Paris"} +FAKE_TOOL = { + "type": "function", + "function": { + "name": TOOL_NAME, + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +# Full parser matrix lives in test_safetensors_tool_loop.py. +TOOL_CALL_TEXT = '{"name": "get_weather", "arguments": {"city": "Paris"}}' +FINAL_ANSWER = "The weather in Paris is sunny and 22C." +TOOL_RESULT = "Paris: sunny, 22C" + + +class RecordingTokenizer: + """Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``. + + Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the + real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually + forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen`` + holding ``None``. + """ + + def __init__(self): + self.tools_seen: list = [] + self.call_count = 0 + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ): + self.call_count += 1 + self.tools_seen.append(kwargs.get("tools")) + return "PROMPT" + + +class StubExecutor: + """Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result. + + A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run. + """ + + def __init__(self, result: str): + self.result = result + self.calls: list[tuple[str, dict]] = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + return self.result + + +def _collect(generator, max_events = 200): + events = [] + for ev in generator: + events.append(ev) + if len(events) >= max_events: + break + return events + + +def _tool_names(tools): + return [(t.get("function") or {}).get("name") for t in (tools or [])] + + +def test_backend_seam_injects_tools_and_drives_full_tool_loop(): + """The shared backend seam forwards tools into the chat template, and the loop parses the call, + dispatches it once, feeds the result back, and re-enters generation for the final answer.""" + tok = RecordingTokenizer() + executor = StubExecutor(TOOL_RESULT) + turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER]) + active_tools_seen: list = [] + conversations_seen: list = [] + + def single_turn(conversation, *, active_tools = None): + # Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots. + active_tools_seen.append(active_tools) + conversations_seen.append([dict(m) for m in conversation]) + apply_chat_template_for_generation(tok, conversation, tools = active_tools) + text = next(turns) + mid = len(text) // 2 + acc = "" + for chunk in (text[:mid], text[mid:]): + acc += chunk + yield acc + + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "What is the weather in Paris?"}], + tools = [FAKE_TOOL], + execute_tool = executor, + max_tool_iterations = 3, + ) + ) + + # 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools). + assert tok.tools_seen, "tokenizer.apply_chat_template was never called" + assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer" + assert TOOL_NAME in _tool_names(tok.tools_seen[0]) + + # 2. Loop offered the tool to the first generation turn. + assert active_tools_seen and active_tools_seen[0] is not None + assert TOOL_NAME in _tool_names(active_tools_seen[0]) + + # 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result. + tool_starts = [e for e in events if e["type"] == "tool_start"] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME + assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls + assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT + + # 6. Final answer streams after the tool result: loop appended it and re-entered generation. + contents = [e for e in events if e["type"] == "content"] + assert contents and FINAL_ANSWER in contents[-1]["text"] + last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end") + last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content") + assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result" + + # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this: + # the fake generation ignores the conversation). + assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call" + final_turn_convo = conversations_seen[1] + assert any( + TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo + ), "tool result was not fed back into the conversation before the final generation turn" + + # 7. Guard: raw tool-call markup never leaked to the client as content. + for e in contents: + assert "" not in e["text"] + assert TOOL_NAME not in e["text"] diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 39fdd151be..7f47140b8d 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -72,10 +72,8 @@ class TestFunctionStyleTrailingText: assert call == {"name": "python", "arguments": {"code": 'print("")'}} def test_closed_function_with_trailing_prose_heal_path(self): - # Regression: the heal / finalize path (allow_incomplete=True) used to fold - # and the trailing prose into the argument and drop - # the prose from visible content. It must now match the strict path -- keep a - # clean argument and leave the trailing prose outside the call span. + # Regression: the heal path (allow_incomplete=True) must match the strict path -- + # keep a clean argument and leave trailing prose outside the call span. text = "cats trailing words" calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 @@ -102,6 +100,24 @@ 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): + # The attribute form (MiniCPM-5 / MiniMax-M2) also ends at the + # LAST , so a literal close tag inside a code argument 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 call with no parameters is a valid zero-argument call; strict + # mode must not treat the empty parameter list as a truncated call. + 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" @@ -184,9 +231,8 @@ class TestHealingPathUnaffected: assert calls[0]["function"]["name"] == "web_search" def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): - # allow_incomplete exists for truncated output; a call that DID close - # must parse identically to strict mode, leaving prose after - # out of the last parameter and out of the removal span. + # A call that DID close must parse identically to strict mode, leaving prose after + # out of the last parameter and the removal span. from core.tool_healing import parse_tool_calls_from_text as parse_with_spans text = "cats trailing" @@ -197,3 +243,1448 @@ class TestHealingPathUnaffected: assert text[span[0] : span[1]] == ( "cats" ) + + def test_wrapperless_fallback_calls_carry_spans(self): + # The wrapperless function-XML fallback must report spans too, so with_spans + # consumers strip exactly the promoted markup (through when closed). + 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 TestGlmStrict: + def test_closed_glm_call_is_accepted(self): + text = ( + "get_weather\n" + "city\nParis\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + + def test_unclosed_glm_call_is_rejected(self): + # No close: truncated, reject with Auto-Heal off. + text = "get_weather\ncity\nParis" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestKimiStrict: + _SB = "<|tool_calls_section_begin|>" + _KB = "<|tool_call_begin|>" + _AB = "<|tool_call_argument_begin|>" + _KE = "<|tool_call_end|>" + _SE = "<|tool_calls_section_end|>" + + def test_full_kimi_call_is_accepted(self): + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + self._SE + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "x" + + def test_kimi_call_without_call_end_is_rejected(self): + # Section closed but the call lacks <|tool_call_end|>: reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._SE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + def test_kimi_without_section_end_is_rejected(self): + # No <|tool_calls_section_end|>: truncated section, reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +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_gemma_wrapperless_deep_nesting_is_linear(self): + # Wrapper-less Gemma ``call:f{a:{a:{...}}}`` deep nesting must parse in linear time (no quadratic re-scan). + import time + + def nested(d): + return "call:f{a:" + "{a:" * d + "x:1" + "}" * d + "}" + + def best_ms(depth): + text = nested(depth) + best = float("inf") + for _ in range(5): + t0 = time.perf_counter() + calls = parse_tool_calls_from_text(text) + best = min(best, time.perf_counter() - t0) + assert calls and json.loads(calls[0]["function"]["arguments"]), "nested args dropped" + return best + + t200 = best_ms(200) + t400 = best_ms(400) + assert t400 < t200 * 3.0, (t200, t400) + + 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 code arg literally containing a <|python_tag|>...call(...) string: the real call is the + # outer "python", not the nested "os" -- the scan stays anchored to the first tag. + 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_glm_open_does_not_parse_spaced_prose_as_tool_name(): + # The GLM NAME opener must reject spaced literal prose (V10); only a + # valid [\w.\-]+ name (followed by newline//) is a call. + assert parse_tool_calls_from_text("not a call") == [] + ok = parse_tool_calls_from_text( + "get_weather\ncity\nNYC\n" + ) + assert [c["function"]["name"] for c in ok] == ["get_weather"] + + +def test_deepseek_r1_missing_call_terminator_rejected_in_strict_mode(): + # R1 must reject a fenced call whose closing ``` + <｜tool▁call▁end｜> never + # arrived when Auto-Heal is off, matching V3/V3.1 strictness (V6). + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}' + "<｜tool▁calls▁end｜>" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +def test_deepseek_r1_complete_call_accepted_in_strict_mode(): + # A fully-terminated R1 call (close fence + per-call end) is still accepted. + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 and calls[0]["function"]["name"] == "get_weather" + + +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_glm_literal_close_tag_in_string_arg_not_truncated(): + import json + + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # A GLM string argument may legitimately contain the literal close tag ````. + text = ( + "run_code\n" + "code\n" + 'print("")\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")', args + + +def test_glm_truncated_block_rejected_in_strict_mode_but_healed_otherwise(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # No close: strict mode (Auto-Heal off) rejects the truncated + # block; with Auto-Heal it keeps the partial call. + text = "get_weather\ncity\nNYC" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 and healed[0]["function"]["name"] == "get_weather" + + +def test_truncated_wrapperless_gemma_call_is_stripped(): + from core.inference.tool_call_parser import strip_tool_markup + + # A wrapper-less Gemma ``call:NAME{...`` cut off mid-arguments (no closing + # brace) must not leak the raw call into the visible stream. + text = 'Sure!\ncall:web_search{"query": "weather in San Fr' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped, repr(stripped) + assert stripped.strip() == "Sure!" + + +def test_complete_wrapperless_gemma_call_keeps_trailing_prose(): + from core.inference.tool_call_parser import strip_tool_markup + + # The truncation pattern must run AFTER the closed form, so a complete call + # followed by prose keeps the prose instead of eating to EOS. + text = 'call:web_search{"query": "cats"} Here you go.' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped + assert stripped.strip() == "Here you go." + + +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 + + # The strip uses the LAST (like the parser) so a literal in a value doesn't + # truncate it; separate calls still 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 parameter value is data, not a call: the scan-based + # strip keeps " done" (the old negative-lookahead regex ate the trailing prose). + 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 emits reasoning as ``[THINK]...[/THINK]`` (bracket form, not ````); + # at end-of-turn it must be dropped so it doesn't leak into display / history. + 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) the reasoning block is left intact; only the + # end-of-turn pass 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 (only a TOP-LEVEL enabled name is a call); the + # ordinary JSON answer is kept verbatim, truncated or complete. + 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 the single-object [TOOL_CALLS]{...} shape, so the display + # strip must remove it too (asymmetry would leak the raw object). + 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(): + # F1: the parser is imported standalone on python >=3.9, where its PEP 604 ``X | None`` + # annotations need ``from __future__ import annotations``; guard that the import 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_glm_strip_treats_literal_close_tag_in_arg_value_as_data(): + # Core strip parity: a literal inside a GLM is argument data, so the whole call is stripped (no leaked tail). + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "web_search\nquery\n" + "see tag\n tail" + ) + assert strip_tool_markup(text, final = True) == "tail" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "see tag"} + + +def test_bare_json_function_alias_parses_and_strips_symmetrically(): + # The bare-JSON parser accepts the "function" alias for the call name; + # strip_leading_bare_json_call must recognise it too (parser/strip symmetry). + 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>", "call:...``), not the wrapper literal anywhere in content: + a query about the marker has nothing tool_healing can parse, and deferring + it loses the call entirely (not executed AND stripped from display).""" + + def test_marker_literal_in_argument_still_parses(self): + text = 'call:web_search{query:"what does <|tool_call> mean"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what does <|tool_call> mean" + + def test_real_wrapped_call_still_deferred_to_tool_healing(self): + from core.inference.tool_call_parser import _parse_gemma_tool_calls + + # An actual wrapped opener present: the Gemma fallback must keep + # deferring to the shared tool_healing parser that owns that form. + text = '<|tool_call>call:web_search{query:<|"|>cats<|"|>}' + assert _parse_gemma_tool_calls(text, id_offset = 0) == [] + + def test_single_quoted_brace_does_not_truncate_code(self): + text = "call:python{code:print('}')}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "print('}')" + + def test_single_quoted_brace_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:python{code:print('}')} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"python"}) + assert "call:python" not in stripped + assert "')}" not in stripped + assert stripped.strip() == "Done." + + +class TestGlmEmbeddedClosePair: + """A GLM value whose string literal embeds the full close-tag pair + ```` (code documenting the GLM format) must not be + truncated at the embedded pair: a structural close sits at balanced quote + state, an embedded one is inside an open string literal.""" + + def test_embedded_pair_inside_quoted_value_not_structural(self): + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")\nx = 1' + + def test_strip_covers_the_full_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + " Done." + ) + stripped = strip_tool_markup(text, final = True) + assert "arg_value" not in stripped + assert stripped.strip() == "Done." + + def test_unbalanced_apostrophe_falls_back_to_first_candidate(self): + # Prose-like value with an apostrophe: no candidate reaches balanced + # quote state, so the first token-valid close wins (prior behavior). + text = ( + "web_search\n" + "query\n" + "it's fine\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "it's fine" + + +class TestPythonTagLiteralInsideMistralArgs: + """A python_tag LITERAL inside a leading Mistral call's arguments is data; the outer call executes.""" + + def test_mistral_arg_quoting_python_tag_call(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "what is <|python_tag|>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 TestGemmaUnquotedApostrophes: + """Quotes open strings only at value-start context: an apostrophe inside + an unquoted wrapper-less value (contractions, possessives) is prose, and + treating it as an opener swallowed the closing brace and lost the call.""" + + def test_contraction_in_unquoted_query_parses(self): + text = "call:web_search{query:what's the weather}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's the weather" + + def test_contraction_does_not_swallow_next_key(self): + text = "call:web_search{query:what's up, n:3}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's up" + assert args["n"] == 3 + + def test_contraction_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:web_search{query:what's the weather} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + assert "call:web_search" not in stripped + assert stripped.strip() == "Done." + + def test_quoted_values_still_hide_delimiters(self): + text = 'call:web_search{query:"weather, location: Boston", n:2}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "weather, location: Boston" + assert args["n"] == 2 + + +class TestGlmKeyWithoutValue: + """A GLM with no tag: strict mode rejects the call + (same contract as an unclosed value) instead of executing it with the + argument silently dropped; Auto-Heal keeps the lenient skip.""" + + def test_strict_rejects_key_without_value(self): + text = "web_search\nquery\n" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_heal_keeps_the_lenient_skip(self): + text = "web_search\nquery\n" + 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"]) == {} + + +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 TestDeepSeekMarkerInsideLeadingEnvelopes: + """A DeepSeek/Kimi marker quoted inside a leading bare-JSON or Mistral + call's argument strings is data: the pre-pass must not promote the + embedded no-arg literal and drop the real outer call.""" + + def test_marker_inside_leading_json_call_stays_data(self): + text = ( + '{"name": "web_search", "arguments": ' + '{"query": "what is <｜tool▁calls▁begin｜>...{}..."}}' + ) + 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 "tool▁calls▁begin" in args["query"] + + def test_marker_inside_leading_mistral_call_stays_data(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "docs on <｜tool▁calls▁begin｜> markers"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_deepseek_call_still_parses(self): + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls = parse_tool_calls_from_text(text) + 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 TestDisabledJsonBeforeDeepSeekCall: + """A disabled leading bare-JSON object whose strings mention a + DeepSeek/Kimi marker is dropped and the tail parsed, so a REAL + DeepSeek/Kimi call after the object still executes instead of the whole + message skipping the pre-pass.""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + + def test_real_deepseek_call_after_disabled_json_parses(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"} ' + self._DS + 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_disabled_json_with_marker_alone_stays_data(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +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 TestLeadingWrapperlessGemmaOverEmbeddedMarkers: + """A leading wrapper-less Gemma call to an enabled tool owns the turn: a + quoted foreign literal inside its argument (a query citing another tool + syntax) is data, and tool_healing must not promote it before the Gemma + fallback runs. Foreign markup leading keeps the normal order.""" + + def test_leading_gemma_wins_over_quoted_xml_literal(self): + text = ( + 'call:web_search{query:"explain ' + '{"name":"evil","arguments":{}}"}' + ) + 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_leading_keeps_normal_order(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"call:evil{x:1} example"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +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_function_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 TestJsonAnswersAreDataForMarkerlessScans: + """A whole-content JSON value is a structured answer: a quoted example of + an enabled tool's syntax inside it must not execute the tool, and the + display strip must not mutilate the answer.""" + + def test_gemma_example_inside_json_answer_not_promoted(self): + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_gemma_example_inside_json_answer_not_stripped(self): + from core.inference.tool_call_parser import strip_tool_markup + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_kimi_marker_inside_json_answer_not_promoted(self): + text = ( + '{"answer":"<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaNestedQuotedLeaves: + def test_nested_object_and_array_values_are_unquoted(self): + text = 'call:f{loc:{city:"New York"},items:["a","b"],n:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New York"}, "items": ["a", "b"], "n": 3} + + +class TestEarliestEnvelopeWinsAcrossDeepSeekKimi: + """The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener: a + leading real call wins over a trailing example of the sibling format in + either direction (document order, like the other leading guards).""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>evil\n" + '```json\n{"x": 1}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + _KIMI = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query": "cats"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_leading_kimi_wins_over_trailing_deepseek_example(self): + text = self._KIMI + " For reference: " + self._DS + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_deepseek_wins_over_trailing_kimi_example(self): + text = self._DS + " Kimi format: " + self._KIMI + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +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 TestClosedCallPrecedesMarkerPrePass: + """A closed non-DeepSeek/Kimi call that precedes the first DS/Kimi marker + owns the turn: a trailing example (or an example quoted inside a wrapped + Gemma argument) must not be promoted by the pre-pass.""" + + _KIMI_EVIL = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.evil:0" + '<|tool_call_argument_begin|>{"x": 1}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_kimi_example_inside_wrapped_gemma_arg_stays_data(self): + text = ( + '<|tool_call>call:web_search{query:<|"|>explain ' + + self._KIMI_EVIL + + '<|"|>}' + ) + 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_leading_xml_call_wins_over_trailing_kimi_example(self): + text = ( + '{"name":"web_search","arguments":{"query":"cats"}}' + " For reference: " + self._KIMI_EVIL + ) + 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_standalone_kimi_call_still_parses(self): + calls = parse_tool_calls_from_text(self._KIMI_EVIL) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestTruncatedWrapperlessGemmaStopsScan: + def test_call_quoted_inside_truncated_arg_not_promoted(self): + text = 'call:python{code:example("call:web_search{query:hi}") and then it cut' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + +class TestGemmaQuotedNestedDelimiters: + def test_comma_inside_quoted_nested_string_not_a_split(self): + text = 'call:f{loc:{city:"New, York"},n:1}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New, York"}, "n": 1} + + +class TestGemmaStringMarkerLiteralInArgs: + def test_string_marker_literal_does_not_lose_the_call(self): + text = "call:web_search{query:'what does <|\"|> mean in Gemma'}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == 'what does <|"|> mean in Gemma' + + +class TestGemmaMidValueQuotedPhrase: + def test_quoted_phrase_mid_value_hides_delimiters(self): + text = 'call:web_search{query:find "weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": 'find "weather, location: Boston"', "limit": 3} + + def test_apostrophes_still_prose_mid_value(self): + text = "call:web_search{query:what's on at the museum, n:2}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's on at the museum", "n": 2} + + +class TestGlmStrictRefusesInQuoteFallback: + """A truncated GLM value whose only close candidates sit inside a string + literal must reject in strict mode instead of executing truncated + arguments; Auto-Heal keeps the lenient partial value.""" + + _TRUNC = ( + 'python\ncode\nprint("")' + ) + + def test_strict_rejects_truncated_in_string_close(self): + assert parse_tool_calls_from_text(self._TRUNC, allow_incomplete = False) == [] + + def test_heal_keeps_partial_value(self): + calls = parse_tool_calls_from_text(self._TRUNC, allow_incomplete = True) + assert len(calls) == 1 and calls[0]["function"]["name"] == "python" + + +class TestGemmaGuardCoversPreambles: + def test_preamble_then_gemma_call_quoting_xml_wins(self): + text = ( + "Sure, searching now. call:web_search{query:" + '"explain {"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGlmStrictAcceptsApostrophes: + def test_apostrophe_value_parses_in_strict_mode(self): + text = ( + "web_search\nquery\n" + "what's the weather\n" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's the weather"} + + +class TestDisabledGemmaCallLiteralsAreData: + def test_literal_inside_disabled_call_not_promoted(self): + text = 'call:foo{query:"x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + def test_real_call_after_disabled_example_still_parses(self): + text = ( + 'call:foo{query:"x"}' + " call:web_search{query:hi}" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingJsonArrayAnswerIsData: + def test_kimi_marker_inside_json_array_answer_not_promoted(self): + text = ( + '[{"answer": "<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}]' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"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 answers and disabled-name objects take the decline path: + # the object is 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 mentioning a literal close tag never folds into arguments. + 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"}}') == "" + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + Gemma span (a quoted plus a later real ).""" + + def test_literal_function_in_gemma_arg_with_later_real_call(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " ls" + " after" + ) + assert strip_tool_call_markup(text, final = True) == "before after" + + def test_literal_function_in_gemma_arg_with_prose_closer(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " then use to close. after" + ) + out = strip_tool_call_markup(text, final = True) + assert out.startswith("before") + assert out.endswith("after") + assert "call:python" not in out + + def test_gemma_opener_inside_json_arg_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + ) + assert strip_tool_call_markup(text, final = True) == "after" + + def test_gemma_opener_inside_function_param_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'x = "<|tool_call>call:t{"' + " after" + ) + assert strip_tool_call_markup(text, final = True) == "after" diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py new file mode 100644 index 0000000000..dfa3101882 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""strip_tool_patterns must match the plain per-pattern loop while skipping the +quadratic no-match rescan of a closed-pair sweep whose close token is absent.""" + +import random +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.tool_healing import ( + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + strip_tool_call_markup, + strip_tool_patterns, +) + + +def _naive(text, patterns): + for pat in patterns: + text = pat.sub("", text) + return text + + +_TOKENS = [ + "", + "", + "<|tool_call>", + "", + "", + "", + "", + "", + "", + "call:fn{", + "}", + "{", + '<|"|>', + "A", + " ", + "\n", + "id", + "x:1", + "", +] + + +def test_guard_matches_plain_loop_on_fuzz(): + rng = random.Random(1234) + for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): + for _ in range(20000): + s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + + +def test_strip_markup_representative_cases_unchanged(): + assert strip_tool_call_markup("a {} b") == "a b" + assert strip_tool_call_markup("a 1 b") == "a b" + # Non-final keeps an unclosed block; final strips it to EOF. + assert strip_tool_call_markup("a {partial") == "a {partial" + assert strip_tool_call_markup("a {partial", final = True) == "a" + + +def test_no_quadratic_blowup_on_unclosed_markers(): + # Unguarded, this took minutes. + big = "" * 20000 + "" * 20000 + t0 = time.perf_counter() + out = strip_tool_call_markup(big, final = True) + assert time.perf_counter() - t0 < 2.0 + assert out == "" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index c2dc1fe8db..d50c27130f 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,18 +24,67 @@ 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} +# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped; +# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here. +assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group( + 1 +), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)" +# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted +# ``_re.compile`` expression resolves the same source. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) + +from typing import Optional as _Optional + +_ns = { + "_re": _re, + "_DS_OPEN_SRC": _DS_OPEN_SRC, + "Optional": _Optional, + "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls, + "_strip_glm_calls": _strip_glm_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"] + +# Signatures may span multiple lines and now carry the enabled_tool_names gate; match +# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. +_xml_helper = _re.search( + r"def _strip_tool_xml\((?:.|\n)*?\) -> 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) +_strip_tool_xml = _ns["_strip_tool_xml"] + _helper = _re.search( - r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" - r"(?: .+\n)+", + r"def _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?: .+\n)+", _src, ) assert _helper, "could not extract _strip_tool_xml_for_display source" +# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the +# extracted body actually reached that call rather than truncating early. +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"] +_gate_src = _re.search( + r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", + _src, +) +assert _gate_src, "could not extract _gemma_strip_gate source" +exec(_gate_src.group(0), _ns) +_gemma_strip_gate = _ns["_gemma_strip_gate"] + # ── Well-formed pairs ───────────────────────────────────────────── @@ -46,6 +95,16 @@ 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_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral + # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). + 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 +132,26 @@ def test_strips_function_only_well_formed(): assert "Done." in cleaned +def test_strips_function_attribute_form(): + # Attribute form ```` (MiniCPM-5 / MiniMax-M2) must strip from the route too + # (it previously leaked into the UI); a dotted/hyphenated name also strips. + 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 + + +# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ── + + +def test_strips_deepseek_space_opener_variant(): + # The space-separated opener is parsed by the parser, so the display strip + # must remove it too (the shared opener alternation is reused here). + text = ( + "pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool" not in cleaned.replace("post", "").replace("pre", "") + assert cleaned == "pre post" + + +def test_strips_deepseek_escaped_underscore_opener_variant(): + text = ( + "pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "pre post" + + +def test_strips_bare_kimi_call_without_section_wrapper(): + # Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no + # section wrapper; the parser accepts it, so the strip must cover it. + text = ( + "pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool_call_begin" not in cleaned + assert cleaned == "pre post" + + +@pytest.mark.parametrize( + "text", + [ + # Prose that merely names a Kimi/DeepSeek marker (no real call follows) must + # survive: the call-shaped lookahead fires only on a real call or a bare EOF + # fragment, so an answer discussing the protocol is never truncated. + "See <|tool_call_begin|> in the docs. More prose after it.", + "The <|tool_calls_section_begin|> marker opens a batch. Read on.", + "DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.", + ], +) +def test_deepseek_kimi_false_alarm_prose_is_kept(text): + # Regression for the route arm truncating a prose answer that references a marker + # without a following call (parser _TOOL_ALL_PATS already had this lookahead). + assert _TOOL_XML_RE.sub("", text) == text + + +def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix(): + # The lookahead must not weaken real-call stripping: closed, truncated, and bare + # EOF-fragment forms all still get removed. + closed = ( + "answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> tail' + ) + assert _TOOL_XML_RE.sub("", closed) == "answer tail" + eof_fragment = "prefix <|tool_call_begin|>" + assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix " + deepseek = ( + "reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + assert _TOOL_XML_RE.sub("", deepseek) == "reply " + + +# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ────── + + +# Llama-3 <|python_tag|> arm bounds on REAL sentinels only +def test_python_tag_strip_consumes_literal_sentinel_in_arg(): + # A <|python_tag|> tool call whose JSON argument carries a literal <|...|> + # token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped + # at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display. + 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 genuine Llama control sentinel still bounds the strip so following + # assistant text is preserved (the arm must not swallow past it). + 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 tool-call region, so the whole pair is + # stripped (the arm bounds the first, then the next match consumes the rest). + 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_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole(): + # GLM 4.x emits NAMEkv .... + text = ( + "web_search\nquery\n" + "find here\n done" + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "" not in out + assert out.strip() == "done" + + +def test_glm_normal_and_qwen_calls_still_stripped_by_route(): + # Regression: a normal GLM call (no literal close tag) and a Qwen + # {json} are still stripped; trailing prose is kept. + glm = "get_time\ntz\nUTC\n ok" + assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok" + qwen = '{"name":"web_search","arguments":{"q":"x"}} after' + assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + + +def test_route_strip_removes_param_alias_close_tag(): + # The parser accepts the ... attribute-form alias of + # ; the route tail cleanup must strip an orphan close 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: the route runs the + # parser's guarded function-XML scan before the regex, matching the core strip. + text = " tail" + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + + +def test_route_strip_gates_wrapperless_gemma_by_enabled_tools(): + # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names, + # like the parser/loop, so a disabled/example name in prose is preserved in ... + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"}) + # An enabled name is still a real call and stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", {"web_search"} + ) + # No gate (legacy) strips every closed call. + assert "call:foo" not in _strip_tool_xml(prose) + + +def test_gemma_strip_gate_empty_tools_preserves_prose(): + # With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls + # back to strip-all and deletes an answer that documents the call:NAME{...} syntax. + assert _gemma_strip_gate([]) == set() + assert _gemma_strip_gate(None) == set() + assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"} + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([])) + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None)) + # An enabled tool's real call is still stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", + _gemma_strip_gate([{"function": {"name": "web_search"}}]), + ) + + +def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): + # The call ends at its first non-data close: prose after it survives the + # strip even when it mentions a literal . + 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 a normal answer must not lose everything after + # it; only text that looks like that family's call start 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(): + # The loops keep this text as next-turn history: a leftover executed call + # would be replayed alongside the structured tool_calls. + 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"}' + ) diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 15961ac03a..2dd2247396 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,341 +1,204 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. +"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Xet (``hf_xet``) is the fast default but can hang with no progress and no -exception, and a blocked native thread cannot be killed. Keep Xet primary; fall -back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` -is read at import time, so the fallback runs in a fresh ``spawn`` child (not a -thread) that sets the env before importing ``huggingface_hub``. Cached files -short-circuit with no child; deterministic errors (401/403/404/disk-full) and -cancellation propagate without a fallback. Mirrors the safetensors inference -recovery in core/inference/{orchestrator,worker}.py. +Re-exports the shared API and injects Studio's marker-aware cache purge +(``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` +marker semantics on the HTTP retry. """ from __future__ import annotations -import multiprocessing as mp -import os -import queue -import signal -import sys import threading -import time from typing import Any, Callable, Optional -from loggers import get_logger +_shared_import_error = None +try: + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True +except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio + # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT + # path before giving up. + _shared_import_error = _exc + import os as _os -logger = get_logger(__name__) - -_CTX = mp.get_context("spawn") - -# Defaults match the existing inference watchdog and hub shutdown deadline. -DEFAULT_HEARTBEAT_INTERVAL = 30.0 -DEFAULT_STALL_TIMEOUT = 180.0 -DEFAULT_GRACE_PERIOD = 10.0 -_POLL_INTERVAL = 0.5 - - -class DownloadStallError(RuntimeError): - """Raised when no download progress is observed for too long. - - Canonical home; orchestrator.py re-imports it so all paths share one type. - """ - - -def child_should_disable_xet(config: dict) -> bool: - """Single source of truth for the per-worker Xet env flip.""" - return bool(config.get("disable_xet")) - - -def get_hf_download_state( - repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" -) -> Optional[tuple[int, bool]]: - """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. - - Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` - is not mistaken for full-size progress. ``None`` means the state could not be - measured, so callers skip stall logic for that tick. - """ + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" try: - from hub.utils.hf_cache_state import ( - blob_bytes_present, - has_active_incomplete_blobs, - hf_cache_root, - iter_active_repo_cache_dirs, - ) + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True + _shared_import_error = None + except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads + _shared_import_error = _exc2 + _shared_available = False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init - if hf_cache_root() is None: - return (0, False) +if _shared_available: + # Bind by assignment so each public name shares one module-level binding with the degraded branch. + DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD + DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL + DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT + DownloadStallError = _shared.DownloadStallError + child_should_disable_xet = _shared.child_should_disable_xet + get_hf_download_state = _shared.get_hf_download_state + start_watchdog = _shared.start_watchdog + _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback + _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback +else: + # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, + # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. + import logging as _logging - total = 0 - has_incomplete = False - for repo_id in repo_ids or []: - # Skip local paths: HF IDs never start with / . ~ or contain "\". - if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: - continue - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - continue - for blob in blobs_dir.iterdir(): - try: - if blob.is_file(): - total += blob_bytes_present(blob) - except OSError: - pass - if has_active_incomplete_blobs(repo_type, repo_id): - has_incomplete = True - return (total, has_incomplete) - except Exception as e: - logger.debug("Failed to determine HF download state: %s", e) - return None + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + DEFAULT_HEARTBEAT_INTERVAL = 30.0 + DEFAULT_STALL_TIMEOUT = 180.0 + DEFAULT_GRACE_PERIOD = 10.0 -def start_watchdog( - *, - repo_ids: list[str], - on_stall: Callable[[str], None], - repo_type: str = "model", - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - stall_timeout: float = DEFAULT_STALL_TIMEOUT, - xet_disabled: bool = False, - on_heartbeat: Optional[Callable[[str], None]] = None, -) -> threading.Event: - """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a - ``*.incomplete`` is present AND the on-disk size is unchanged for - *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so - post-download init is never misread as a stall. Returns a stop event the - caller sets when the download phase ends. - """ - stop = threading.Event() - transport = "https" if xet_disabled else "xet" - fired = False + class DownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" - def _beat() -> None: - nonlocal fired - state = get_hf_download_state(repo_ids, repo_type = repo_type) - last_size = state[0] if state is not None else 0 - last_change = time.monotonic() + def child_should_disable_xet(config: dict) -> bool: + return bool(config.get("disable_xet")) - while not stop.wait(interval): - state = get_hf_download_state(repo_ids, repo_type = repo_type) - now = time.monotonic() + def get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires - if state is None: - if on_heartbeat is not None: + def start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, + ) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: + return stop + transport = "https" if xet_disabled else "xet" + + def _beat() -> None: + while not stop.wait(interval): + try: on_heartbeat(f"Downloading ({transport} transport)...") - continue + except Exception: + pass - current_size, has_incomplete = state - if current_size != last_size: - last_size = current_size - last_change = now + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - # Reset unless .incomplete confirms an active download, so model init - # and lock waits are not counted as a stall. - if not has_incomplete: - last_change = now - elif now - last_change >= stall_timeout: - if not fired: - fired = True - on_stall( - f"Download appears stalled ({transport} transport) " - f"-- no progress for {int(now - last_change)}s" - ) - return + def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - if on_heartbeat is not None: - on_heartbeat(f"Downloading ({transport} transport)...") + def _shared_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() - return stop - - -def _download_child_entry( - *, - repo_id: str, - filename: str, - token: Optional[str], - repo_type: str, - disable_xet: bool, - result_queue: Any, - force_download: bool = False, -) -> None: - """Spawn-child entrypoint: download one file and report the result. - - Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, - forms its own process group so the parent can kill the whole transfer, and - never logs the token or signed URLs. - """ - # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). - try: - from utils.process_lifetime import bind_current_process_to_parent_lifetime - bind_current_process_to_parent_lifetime() - except Exception: - pass - - if hasattr(os, "setsid"): - try: - os.setsid() - except OSError: - pass - - if disable_xet: - os.environ["HF_HUB_DISABLE_XET"] = "1" - # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse - # partials a sequential resume cannot safely continue). - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - # Test-only fault injection (never set in production): stall the Xet attempt - # so the watchdog + HTTP fallback can be exercised against a real repo. - if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": - import time as _t - try: - from huggingface_hub.constants import HF_HUB_CACHE - - blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") - os.makedirs(blobs, exist_ok = True) - with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: - fh.write(b"\0" * 4096) - except OSError: - pass - while True: - _t.sleep(3600) - - try: from huggingface_hub import hf_hub_download + path = hf_hub_download( repo_id = repo_id, filename = filename, - repo_type = repo_type, token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, force_download = force_download, ) - result_queue.put({"ok": True, "path": path}) - except BaseException as e: # noqa: BLE001 - report every failure to the parent - error = f"{type(e).__name__}: {e}" - try: - from hub.utils.download_registry import scrub_secrets - error = scrub_secrets(error, hf_token = token) - except Exception: - pass - result_queue.put({"ok": False, "error": error}) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + def _shared_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") -def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: - """Kill *proc* and its whole process group (Xet may spawn helper procs). + from huggingface_hub import snapshot_download - The child calls ``os.setsid()`` so its pgid equals its pid; signal via - ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a - group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. - """ - pid = proc.pid - - def _signal_group(sig: int) -> None: - if pid is not None and hasattr(os, "killpg"): - try: - os.killpg(pid, sig) - return - except (ProcessLookupError, PermissionError, OSError): - pass - # Windows or pre-setsid: best effort on the single process. - try: - proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() - except Exception: - pass - - _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) - proc.join(timeout = grace_period) - if proc.is_alive(): - _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) - proc.join(timeout = 5.0) - - -def _run_download_attempt( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str, - disable_xet: bool, - cancel_event: Optional[threading.Event], - stall_timeout: float, - interval: float, - grace_period: float, - on_status: Optional[Callable[[str], None]], - force_download: bool = False, -) -> tuple[str, Optional[str]]: - """Run one download in a spawn child supervised by the no-progress watchdog. - - Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or - ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. - """ - result_queue: Any = _CTX.Queue() - proc = _CTX.Process( - target = _download_child_entry, - kwargs = dict( + path = snapshot_download( repo_id = repo_id, - filename = filename, - token = token, repo_type = repo_type, - disable_xet = disable_xet, - result_queue = result_queue, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, force_download = force_download, - ), - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid - - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) - - stalled = threading.Event() - stop_watchdog = start_watchdog( - repo_ids = [repo_id], - on_stall = lambda msg: stalled.set(), - repo_type = repo_type, - interval = interval, - stall_timeout = stall_timeout, - xet_disabled = disable_xet, - on_heartbeat = on_status, - ) - - result: Optional[dict] = None - try: - while proc.is_alive(): - if cancel_event is not None and cancel_event.is_set(): - _terminate_process_group(proc, grace_period) - return ("cancelled", None) - if stalled.is_set(): - _terminate_process_group(proc, grace_period) - return ("stall", None) - try: - result = result_queue.get(timeout = _POLL_INTERVAL) - break - except queue.Empty: - continue - else: - # Process exited; drain any result it enqueued. - try: - result = result_queue.get_nowait() - except queue.Empty: - result = None - finally: - stop_watchdog.set() - proc.join(timeout = grace_period) - - if result is None: - return ( - "error", - f"download process for '{repo_id}/{filename}' exited " - f"(code={proc.exitcode}) without a result", ) - if result.get("ok"): - return ("ok", result["path"]) - return ("error", result.get("error") or "unknown download error") + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +__all__ = [ + "DEFAULT_GRACE_PERIOD", + "DEFAULT_HEARTBEAT_INTERVAL", + "DEFAULT_STALL_TIMEOUT", + "DownloadStallError", + "child_should_disable_xet", + "get_hf_download_state", + "start_watchdog", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", +] + + +def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: + """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, + not fatal to the retry.""" + try: + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as exc: + try: + from loggers import get_logger + get_logger(__name__).debug( + "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + ) + except ModuleNotFoundError as logger_exc: + if logger_exc.name != "loggers": + raise def hf_hub_download_with_xet_fallback( @@ -345,83 +208,32 @@ def hf_hub_download_with_xet_fallback( *, cancel_event: Optional[threading.Event] = None, repo_type: str = "model", + revision: Optional[str] = None, stall_timeout: float = DEFAULT_STALL_TIMEOUT, interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, ) -> str: - """Download a single file with Xet primary and HTTP as a stall-only fallback. + """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + return _shared_hf_hub_download_with_xet_fallback( + repo_id, + filename, + token, + cancel_event = cancel_event, + repo_type = repo_type, + revision = revision, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + force_download = force_download, + prepare_for_http_fn = _studio_prepare_for_http, + ) - Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if - *cancel_event* is set, re-raises a deterministic child error unchanged (no - fallback), and raises ``DownloadStallError`` only if BOTH transports stall. - When *force_download* is True the cache-first early-return is skipped and the - flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched - even if an older blob is already cached. - """ - # Finalized blob already cached: return it with no child and no network. - # Skipped when force_download is set so an update re-fetches a newer blob. - if not force_download: - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) - - if cancel_event is not None and cancel_event.is_set(): - raise RuntimeError("Cancelled") - - disable_xet = False - for attempt in range(2): - if disable_xet: - # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume - # over a sparse Xet/hf_transfer partial silently corrupts the blob. - try: - from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") - except Exception as e: - logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) - - kind, payload = _run_download_attempt( - repo_id, - filename, - token, - repo_type = repo_type, - disable_xet = disable_xet, - cancel_event = cancel_event, - stall_timeout = stall_timeout, - interval = interval, - grace_period = grace_period, - on_status = on_status, - force_download = force_download, - ) - - if kind == "ok": - return payload # type: ignore[return-value] - if kind == "cancelled": - raise RuntimeError("Cancelled") - if kind == "error": - # Deterministic failure: the other transport would fail identically. - raise RuntimeError(payload) - # kind == "stall" - if attempt == 0 and not disable_xet: - logger.warning( - "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", - repo_id, - filename, - ) - if on_status is not None: - on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") - disable_xet = True - continue - raise DownloadStallError( - f"Download stalled for '{repo_id}/{filename}' even with " - f"HF_HUB_DISABLE_XET=1 -- check your network connection" - ) - - # Unreachable: the loop either returns or raises on each attempt. - raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") +def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: + """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx index d00c812325..db7628777a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -42,10 +42,8 @@ export function ModelUpdateAction({ }: ModelUpdateActionProps) { const [open, setOpen] = useState(false); - // The update is a managed download (it surfaces in the global Downloads panel - // with progress + cancel). When this exact repo+variant finishes, refresh the - // caller so the "update available" cue clears once the new revision is on - // disk. A ref keeps the subscription stable across renders without resubscribing. + // Refresh the caller when this repo+variant's download finishes so the "update available" cue + // clears. A ref keeps the subscription stable across renders. const onUpdatedRef = useRef(onUpdated); onUpdatedRef.current = onUpdated; useEffect(() => { @@ -60,9 +58,8 @@ export function ModelUpdateAction({ }, [repoId, variant]); const handleConfirm = useCallback(() => { - // Start the background re-download and close the dialog immediately; the - // Downloads panel owns progress + cancel from here. Only a failure to START - // surfaces a toast — a failed download reports itself in the panel. + // Start the re-download and close the dialog; the Downloads panel owns progress + cancel. + // Only a failure to START toasts (a failed download shows in the panel). void Promise.resolve() .then(onConfirm) .catch((err) => { diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 5af37ee192..e786cf8ac0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1505,11 +1505,8 @@ export function HubModelPicker({ catalog?: CatalogGroup[]; }) { const gpu = useGpuInfo(); - // The currently-loaded/running model id. We read params.checkpoint from the - // runtime store (backend-mirrored from /api/inference/status.active_model, see - // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is - // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used - // to disable the cached-row update action for the model that's live in memory. + // Live model id from the runtime store (backend-mirrored active_model), not the dropdown + // highlight which can be a staged pick. Disables the update action for it. const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). const loadTimes = useModelLoadTimes(value); @@ -1849,11 +1846,8 @@ export function HubModelPicker({ refreshLocalModelsList(); }, [hfToken, refreshLocalModelsList]); - // Updates run as MANAGED downloads (they show in the global Downloads panel - // with manifest-based progress + a working Cancel), instead of a blocking - // call. The worker re-resolves `main` and pulls only changed blobs, so the - // cached copy stays usable until the new revision lands. The row's - // ModelUpdateAction refreshes the list when this repo+variant completes. + // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking + // call. The worker pulls only changed blobs, so the cached copy stays usable until done. const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { return downloadManager .requestStart({ diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index f02da6612e..bce4bf2831 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -27,9 +27,11 @@ function usageTextClass(percent: number): string { return "text-primary"; } -function formatGb(value: number): string { +function formatGiB(value: number): string { + // RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching + // nvidia-smi and PyTorch, so label the readout GiB rather than GB. const digits = value >= 10 ? 1 : 2; - return `${value.toFixed(digits)} GB`; + return `${value.toFixed(digits)} GiB`; } export function FloatingMonitor() { @@ -116,7 +118,7 @@ export function FloatingMonitor() {
- {formatGb(ramUsed)} / {formatGb(ramTotal)} + {formatGiB(ramUsed)} / {formatGiB(ramTotal)}
- {formatGb(vramUsed)} / {formatGb(vramTotal)} + {formatGiB(vramUsed)} / {formatGiB(vramTotal)}
0 - ? `${Math.round(gpu.systemRamTotalGb)} GB` + ? `${Math.round(gpu.systemRamTotalGb)} GiB` : "Unavailable"; const coreLabel = gpu.cpuCore > 0 && gpu.cpuThread > 0 diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 6b6b11bf1f..eb50d398b6 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -125,7 +125,7 @@ export function SummaryStep() { GPU
{hw.gpuName ?? "---"} - {hw.vramTotalGb != null ? `${hw.vramTotalGb} GB` : "---"} + {hw.vramTotalGb != null ? `${hw.vramTotalGb} GiB` : "---"}
diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx index ff751e3cd6..1f323a9e8d 100644 --- a/studio/frontend/src/features/settings/tabs/about-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx @@ -158,7 +158,7 @@ export function AboutTab() { {gpu.name ?? "—"} {gpu.vramTotalGb != null - ? ` · ${Math.round(gpu.vramTotalGb)} GB` + ? ` · ${Math.round(gpu.vramTotalGb)} GiB` : ""} diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx index d5e19cc51c..6c30858c63 100644 --- a/studio/frontend/src/features/settings/tabs/resources-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -47,6 +47,15 @@ function formatGb(value: number | null | undefined): string { return `${safe.toFixed(digits)} GB`; } +// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching +// nvidia-smi and PyTorch, so label those readouts GiB. Disk stays on formatGb +// because the backend reports disk in decimal GB (bytes / 1e9). +function formatGiB(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GiB`; +} + function formatMb(value: number | null | undefined): string { const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; return `${Math.round(safe).toLocaleString()} MB`; @@ -300,9 +309,9 @@ export function ResourcesTab() { /> @@ -318,13 +327,13 @@ export function ResourcesTab() { label={t("settings.resources.liveMonitor.vram")} value={ hasGpu - ? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}` + ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}` : t("settings.resources.liveMonitor.noGpu") } detail={ hasGpu ? t("settings.resources.liveMonitor.free", { - value: formatGb(metrics.vramFree), + value: formatGiB(metrics.vramFree), }) : backendLabel } @@ -373,17 +382,17 @@ export function ResourcesTab() {
{t("settings.resources.gpu.used", { - value: formatGb(used), + value: formatGiB(used), })} {t("settings.resources.gpu.free", { - value: formatGb(free), + value: formatGiB(free), })} {t("settings.resources.gpu.total", { - value: formatGb(total), + value: formatGiB(total), })}
diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index abab35db93..9c9398688f 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -411,7 +411,7 @@ function LiveGpuPanel({ value={index} className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100" > - GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"}) + GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GiB` : "N/A"}) ))} @@ -446,7 +446,7 @@ function LiveGpuPanel({ icon={} value={ currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null - ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GiB` : "--" } pct={currentGpu.vram_utilization_pct ?? 0} diff --git a/tests/saving/test_quant_method_none_normalization.py b/tests/saving/test_quant_method_none_normalization.py new file mode 100644 index 0000000000..c1c5fd3686 --- /dev/null +++ b/tests/saving/test_quant_method_none_normalization.py @@ -0,0 +1,81 @@ +"""CPU-only regression for the quant-method normalization loops in save.py. + +`unsloth_save_pretrained_gguf` and `save_to_gguf_generic` each normalize the +`quantization_method` list, mapping a ``None`` element to ``"q8_0"``. The mapping +used to call ``quant_method.lower()`` as the first statement of the loop, so a +``None`` element (e.g. ``quantization_method=[None]`` or ``["q4_k_m", None]``) +raised ``AttributeError: 'NoneType' object has no attribute 'lower'`` and the +``elif quant_method is None`` branch was unreachable dead code. + +The loop is inline inside two heavy functions (importing unsloth needs +unsloth_zoo / a GPU), so - like test_is_gpt_oss_detection.py - we extract just the +loop source via ``ast`` and exec it against sample inputs. That exercises the real +source: it fails on the old ordering and passes once ``None`` is handled first. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" +SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8") +SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY)) + +# The target functions and the list variable each one appends the normalized method to. +TARGETS = ( + ("unsloth_save_pretrained_gguf", "quantization_methods"), + ("save_to_gguf_generic", "new_quantization_methods"), +) + + +def _func(tree, name): + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"function {name!r} not found in {SAVE_PY.name}") + + +def _quant_loop(func_name): + # The quant-normalization `for` loop iterates `quantization_method`; grab its source. + func = _func(SAVE_TREE, func_name) + for node in ast.walk(func): + if ( + isinstance(node, ast.For) + and isinstance(node.iter, ast.Call) + and isinstance(node.iter.func, ast.Name) + and node.iter.func.id == "enumerate" + and isinstance(node.iter.args[0], ast.Name) + and node.iter.args[0].id == "quantization_method" + ): + return node + raise AssertionError(f"quant-normalization loop not found in {func_name}") + + +def _run_loop(func_name, out_var, quantization_method): + # exec just the extracted loop against a given input, returning the appended methods. + loop_src = ast.get_source_segment(SAVE_SRC, _quant_loop(func_name)) + namespace = {out_var: [], "quantization_method": quantization_method} + exec(loop_src, {"__builtins__": __builtins__}, namespace) + return namespace[out_var] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_none_element_maps_to_q8_0(func_name, out_var): + # A bare None inside the list must map to q8_0, not raise AttributeError. + assert _run_loop(func_name, out_var, [None]) == ["q8_0"] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_none_mixed_with_strings(func_name, out_var): + # None resolves to q8_0 while sibling string methods are still normalized (lowercased). + assert _run_loop(func_name, out_var, ["Q4_K_M", None]) == ["q4_k_m", "q8_0"] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_string_methods_unchanged(func_name, out_var): + # The fix must not alter behavior for the ordinary string inputs. + methods = ["not_quantized", "fast_quantized", "quantized", "Q8_0"] + assert _run_loop(func_name, out_var, methods) == ["f16", "q8_0", "q4_k_m", "q8_0"] diff --git a/tests/saving/test_torchao_remote_code_consent.py b/tests/saving/test_torchao_remote_code_consent.py new file mode 100644 index 0000000000..1d0acb7866 --- /dev/null +++ b/tests/saving/test_torchao_remote_code_consent.py @@ -0,0 +1,150 @@ +"""Regression tests for the export-time remote-code trust decision. + +FP8/FP4/INT quantization export re-reads the just-merged checkpoint. It used to enable +trust_remote_code whenever the checkpoint's config carried an ``auto_map`` entry, so a model +that loads fine with built-in classes (and therefore skips the load-time consent scan) could +smuggle unvetted remote code that then runs at export. The export paths now derive +trust_remote_code from ``_loaded_via_remote_code`` - the already approved load decision - instead. + +These run on CPU with no torch / unsloth import: they AST-extract the real helper from +unsloth/save.py and exec it in isolation, plus assert the call sites dropped the auto_map trust. +""" + +import ast +from pathlib import Path + +_SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" +_SRC = _SAVE_PY.read_text(encoding = "utf-8") + + +def _load_helper(): + """Exec just `_loaded_via_remote_code` from save.py (no torch import) and return it.""" + tree = ast.parse(_SRC) + fn = next( + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_loaded_via_remote_code" + ) + ns = {} + exec(compile(ast.Module(body = [fn], type_ignores = []), str(_SAVE_PY), "exec"), ns) + return ns["_loaded_via_remote_code"] + + +_loaded_via_remote_code = _load_helper() + + +def _obj(module_name, **attrs): + """A throwaway instance whose class __module__ is `module_name`, plus given attributes.""" + cls = type("Fake", (), {}) + cls.__module__ = module_name + inst = cls() + for k, v in attrs.items(): + setattr(inst, k, v) + return inst + + +def test_builtin_class_is_not_remote_code(): + assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False + + +def test_transformers_modules_class_is_remote_code(): + assert _loaded_via_remote_code(_obj("transformers_modules.acme.modeling_x")) is True + + +def test_none_is_not_remote_code(): + assert _loaded_via_remote_code(None) is False + + +def test_none_module_is_not_remote_code(): + # A class whose __module__ is None must not raise AttributeError. + assert _loaded_via_remote_code(_obj(None)) is False + + +def test_auto_map_in_config_alone_does_not_grant_trust(): + # The core bypass: a built-in-loadable model whose config merely declares auto_map must NOT + # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass). + cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})() + assert ( + _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg)) + is False + ) + + +def test_peft_base_model_is_unwrapped(): + base = _obj("transformers_modules.acme.modeling_x") + peft = _obj("peft.peft_model", get_base_model = lambda: base) + assert _loaded_via_remote_code(peft) is True + + +def test_wrapper_model_attr_is_walked(): + inner = _obj("transformers_modules.acme.modeling_x") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is True + + +def test_wrapper_over_builtin_stays_false(): + inner = _obj("transformers.models.llama.modeling_llama") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is False + + +def test_processor_held_custom_tokenizer_is_detected(): + # A built-in ProcessorMixin can hold an approved custom-code tokenizer; the walk must + # descend into processor components or the export reload loses that approved trust. + tok = _obj("transformers_modules.acme.tokenization_x") + proc = _obj("transformers.processing_utils", tokenizer = tok) + assert _loaded_via_remote_code(proc) is True + + +def test_processor_held_custom_image_processor_is_detected(): + ip = _obj("transformers_modules.acme.image_processing_x") + proc = _obj("transformers.processing_utils", image_processor = ip) + assert _loaded_via_remote_code(proc) is True + + +def test_builtin_processor_with_builtin_components_stays_false(): + proc = _obj( + "transformers.processing_utils", + tokenizer = _obj("transformers.tokenization_utils_fast"), + image_processor = _obj("transformers.image_processing_utils"), + ) + assert _loaded_via_remote_code(proc) is False + + +def test_cyclic_wrappers_terminate(): + a = _obj("peft.peft_model") + b = _obj("peft.peft_model", model = a) + a.model = b + assert _loaded_via_remote_code(a) is False + + +# -- call-site assertions: the auto_map-derived trust is gone from every export path ----------- + + +def test_torchao_export_derives_trust_from_load_decision(): + assert "model_trust = _loaded_via_remote_code(model)" in _SRC + assert "tok_trust = _loaded_via_remote_code(tokenizer)" in _SRC + assert "trust_remote_code = model_trust" in _SRC + assert "trust_remote_code = tok_trust" in _SRC + # The staged-config auto_map scan that granted trust is removed. + assert 'if "auto_map" in json.load' not in _SRC + + +def test_compressed_and_gguf_lora_paths_drop_auto_map_trust(): + # No path derives a trust decision straight from config auto_map anymore, and no path + # collapses model and tokenizer trust into one flag. + assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC + assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC + assert "if _loaded_via_remote_code(model):" in _SRC # GGUF-LoRA converter flag + + +def test_compressed_export_keeps_model_and_tokenizer_trust_separate(): + # The subprocess gets one flag per component, so an approved custom tokenizer cannot + # enable an unapproved model's code during compressed quantization (or vice versa). + assert 'cmd.append("--trust-remote-code")' in _SRC + assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC + qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8") + assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc + assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc + # The model loads keep the model flag only. + assert "args.model, args.trust_remote_code)" in qsrc diff --git a/tests/test_attn_impl_honor_explicit.py b/tests/test_attn_impl_honor_explicit.py new file mode 100644 index 0000000000..3fb7a2208f --- /dev/null +++ b/tests/test_attn_impl_honor_explicit.py @@ -0,0 +1,190 @@ +"""An explicit non-flash attention request must survive the flash disable path. + +When flash attention is disabled for a model, a caller who explicitly asked for +"sdpa" or "flex_attention" should keep that choice instead of being downgraded +to whatever the conservative supports_* fallback would pick. +""" + +import pytest + +from unsloth.models._utils import ( + _disable_flash_attention_if_needed, + resolve_attention_implementation, +) + + +def test_explicit_sdpa_is_honored_even_when_not_marked_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "sdpa", + supports_sdpa = False, # conservative flag would have skipped sdpa + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_explicit_flex_is_honored_when_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + assert config.get("_attn_implementation") == "flex_attention" + + +def test_explicit_flex_falls_back_when_not_supported(): + # flex_attention is False for known-broken/excluded configs (e.g. gpt_oss), + # so an explicit flex request must not select that backend - it falls back. + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_synthesized_config_sdpa_is_not_treated_as_explicit(): + # The language loader seeds the config with attn_implementation="sdpa"; when the + # caller passes nothing, that synthesized value must not override the flex fallback + # for a model that supports flex but not sdpa. + config = {"attn_implementation": "sdpa"} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = None, + supports_sdpa = False, + supports_flex_attention = True, + would_use_flash_attention = False, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + + +def test_no_disable_reason_returns_request_untouched(): + result = _disable_flash_attention_if_needed( + {}, + attn_implementation = "flash_attention_2", + disable_reason = None, + ) + assert result == "flash_attention_2" + + +def test_flash_request_still_falls_back_when_disabled(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flash_attention_2", + supports_sdpa = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_resolver_honors_explicit_sdpa_when_not_supported_and_flash_disabled(): + # End-to-end through the public resolver: an explicit sdpa request with a + # flash-disabled config (oversized head dim) and supports_sdpa=False must not be + # rewritten to eager by the resolver's own not-supports_sdpa guard. + config = {"model_type": "test", "head_dim": 512} # head_dim > 256 disables flash + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_non_explicit_sdpa_when_not_supported(): + # No explicit request: the model resolution seeds sdpa/eager and the guard must + # still downgrade a synthesized sdpa to eager for a model that cannot run it. + config = {"model_type": "test", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +def test_resolver_downgrades_explicit_sdpa_for_sdpa_excluded_model(): + # gpt_oss is in _SDPA_EXCLUDED_MODELS (sdpa is known-broken) and _FLASH_EXCLUDED_MODELS + # (flash disabled). Honoring an explicit sdpa request must not re-enable that broken + # backend: it downgrades to eager, mirroring how an explicit flex request falls back + # for _FLEX_EXCLUDED_MODELS. supports_sdpa=True proves the exclusion overrides even a + # model that otherwise advertises SDPA support. + config = {"model_type": "gpt_oss"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = True, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +@pytest.mark.parametrize("model_type", ["gemma3", "gemma3_text"]) +def test_resolver_downgrades_explicit_sdpa_for_disable_sdpa_model(model_type): + # gemma3 / gemma3_text are in DISABLE_SDPA_MODEL_NAMES: the loader forces + # supports_sdpa=False because their bundled SDPA modules are wrong. An explicit + # sdpa request with flash disabled must NOT re-enable that known-wrong path - it + # downgrades to eager, exactly like _SDPA_EXCLUDED_MODELS (gpt_oss). head_dim>256 + # disables flash to mirror the real flash-disabled scenario. + config = {"model_type": model_type, "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +def test_resolver_does_not_overmatch_gemma3n_for_explicit_sdpa(): + # The "gemma3," trailing-comma guard must not match gemma3n: gemma3n is not in + # DISABLE_SDPA_MODEL_NAMES, so it stays a conservative (not known-wrong) model and an + # explicit sdpa request is still honored. Proves the substring match neither over- nor + # under-matches. + config = {"model_type": "gemma3n", "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_synthesized_sdpa_for_disable_sdpa_model(): + # A synthesized/default sdpa (requested is None; the value came from config) on a + # DISABLE_SDPA_MODEL_NAMES model must still downgrade to eager. + config = {"model_type": "gemma3", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, "-q"])) diff --git a/tests/test_fp8_tiny_e8m0.py b/tests/test_fp8_tiny_e8m0.py new file mode 100644 index 0000000000..cf49c8c92f --- /dev/null +++ b/tests/test_fp8_tiny_e8m0.py @@ -0,0 +1,123 @@ +"""FP8 block-quant linear must handle tiny / non-tileable weights and e8m0 scales. + +Two things break the triton block path: + * a hidden dim not divisible by the activation block size (tiny test models), + * float8_e8m0fnu weight scales, which have no triton dtype mapping. +The forward falls back to a torch-native blockwise dequant + bf16 matmul; this +test checks that fallback runs finite forward + backward and matches a plain +dequant reference. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA") + + +def _reference(X, weight, scale, block): + # Expand the per-block scale to full weight shape and dequantize. + m, n = weight.shape + s = scale.to(torch.float32) + s = s.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n] + W = (weight.to(torch.float32) * s).to(X.dtype) + return X @ W.T + + +def test_tiny_non_tileable_forward_backward_matches_reference(): + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + torch.manual_seed(0) + dev = "cuda" + block = [128, 128] + m, n = 8, 8 # non-tileable, in-dim % 128 != 0 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) # (out=m, in=n) + scale = torch.rand(1, 1, device = dev, dtype = torch.float32) + 0.5 + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all(), "forward produced non-finite values" + + ref = _reference(X.detach(), weight, scale, block) + torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2) + + out.sum().backward() + assert X.grad is not None and torch.isfinite(X.grad).all(), "backward non-finite" + + +def test_e8m0_scale_is_upcast_and_runs(): + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + if not hasattr(torch, "float8_e8m0fnu"): + pytest.skip("torch build lacks float8_e8m0fnu") + + dev = "cuda" + m, n = 8, 8 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) + scale = (torch.rand(1, 1, device = dev) + 1.0).to(torch.float8_e8m0fnu) + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all() + out.sum().backward() + assert torch.isfinite(X.grad).all() + + +def test_rectangular_block_dequant_matches_reference(): + # Rectangular blocks (block_size[0] != block_size[1]) that tile evenly used to + # route through the triton weight_dequant kernel, which uses a single BLOCK_SIZE + # for both axes and mis-indexes the column scale. Verify the torch expansion path + # now matches the reference for a 64x256 weight with block [64, 128] (scale 1x2). + from unsloth.kernels.fp8 import _blockwise_weight_dequant_any_shape + + torch.manual_seed(0) + dev = "cuda" + block = [64, 128] + m, n = 64, 256 # evenly tiled: 64 % 64 == 0, 256 % 128 == 0 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) + # Distinct per-block column scales expose column mis-indexing. + scale = torch.tensor([[0.5, 3.0]], device = dev, dtype = torch.float32) + + W_deq = _blockwise_weight_dequant_any_shape(weight, scale, block, torch.bfloat16) + + s = scale.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n] + ref = (weight.to(torch.float32) * s).to(torch.bfloat16) + torch.testing.assert_close(W_deq, ref, atol = 5e-3, rtol = 5e-3) + + +def test_e8m0_scale_preserves_non_default_block_size_attr(): + # An e8m0 scale carrying a non-default block_size attribute must keep it across + # the float32 upcast in forward; otherwise the lookup falls back to [128, 128] + # and a compatible layout is wrongly rejected as incompatible. + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + if not hasattr(torch, "float8_e8m0fnu"): + pytest.skip("torch build lacks float8_e8m0fnu") + + torch.manual_seed(0) + dev = "cuda" + block = [64, 64] + # in-dim 96 is not divisible by block[1]=64 -> forward takes the torch dequant + # fallback (no fp8 matmul kernel). Scale shape (2, 2) validates for [64, 64] but + # not [128, 128] (which expects (1, 1)). + m, n = 128, 96 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) # no block_size attr + scale_f = torch.rand(2, 2, device = dev) + 1.0 + scale = scale_f.to(torch.float8_e8m0fnu) + scale.block_size = block # attribute lives on the scale, not the weight + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + # With [128, 128] this raises "not compatible with block size"; success proves + # the [64, 64] attribute survived the e8m0 -> float32 upcast. + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all() + + ref = _reference(X.detach(), weight, scale.to(torch.float32), block) + torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2) + + out.sum().backward() + assert X.grad is not None and torch.isfinite(X.grad).all() + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, "-q"])) diff --git a/tests/test_moe_lora_targets.py b/tests/test_moe_lora_targets.py index 994d39f261..7f9b9a0485 100644 --- a/tests/test_moe_lora_targets.py +++ b/tests/test_moe_lora_targets.py @@ -49,3 +49,190 @@ def test_explicit_dotted_module_target_does_not_discover_moe_parameters(): ) is None ) + + +@pytest.mark.parametrize( + "target_modules", + [ + # Attention-only auto-regex lists every projection leaf (incl. gate/up/down) + # but its path segment is attention-only, so experts must NOT be targeted. + r"(?:\bmodel\.layers\.[\d]{1,}\.(?:self_attn|attention|attn|mixer)\.(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))", + ".*self_attn.*proj", + # An mlp path alternative with attention-only leaves is still attention-only. + r"model\.layers\.\d+\.(?:mlp|self_attn)\.(?:q_proj|k_proj|v_proj|o_proj)", + ], +) +def test_attention_only_regex_does_not_discover_moe_parameters(target_modules): + from unsloth.models._utils import get_moe_target_parameters + assert get_moe_target_parameters(_FakeMoeModel(), target_modules) is None + + +def test_single_leaf_regex_targets_only_that_projection(): + from unsloth.models._utils import get_moe_target_parameters + assert get_moe_target_parameters(_FakeMoeModel(), ".*experts.*down_proj") == [ + "mlp.experts.down_proj", + ] + assert get_moe_target_parameters(_FakeMoeModel(), ".*mlp.*gate_proj") == [ + "mlp.experts.gate_up_proj", + ] + + +def test_auto_regex_mlp_tag_block_discovers_moe_on_fused_models(): + # get_peft_regex on a fused-expert model lists only attention Linears as + # leaves; the mlp tag block is the remaining signal of MLP finetune intent. + from unsloth.models._utils import get_moe_target_parameters + both_auto = ( + r"(?:\bmodel\.layers\.[\d]{1,}\." + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\." + r"(?:(?:q_proj|k_proj|v_proj|o_proj)))" + ) + assert get_moe_target_parameters(_FakeMoeModel(), both_auto) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_explicit_attention_only_list_does_not_discover_moe_parameters(): + # An explicit attention-only leaf list names no MLP projection, so experts + # must never be targeted. get_peft_model routes this ORIGINAL list (not the + # scoped regex) into detection precisely because family scoping makes + # get_peft_regex emit its full "mlp|feed_forward|ffn|dense" component block + # even for an attention-only request (see the regex below), which the + # string fallback cannot distinguish from the fused-expert auto regex. + from unsloth.models._utils import get_moe_target_parameters + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + assert get_moe_target_parameters(_FakeMoeModel(), attn_only_list) is None + assert get_moe_target_parameters(_FakeMoeModel(), tuple(attn_only_list)) is None + + # The regex get_peft_regex emits for that same attention-only list under a + # vision-off family scope carries the mlp component block, so the string + # path would wrongly enable experts -- hence detection must use the list. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + assert get_moe_target_parameters(_FakeMoeModel(), scoped_regex) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_frozen_mlp_full_list_does_not_discover_moe_parameters(): + # Regression: an explicit list that names MLP leaves together with + # finetune_mlp_modules=False must NOT train experts. get_peft_regex scopes + # the MLP leaves out (its emitted regex carries no mlp tag block), so + # detection has to key on that SCOPED regex -- keying on the original list + # would let its gate/up/down leaves silently re-enable the frozen experts. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + # Representative of what get_peft_regex emits for that list under + # finetune_mlp_modules=False: attention-only path, no mlp component block. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = False, + finetune_language_layers = True, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_frozen_language_full_list_does_not_discover_moe_parameters(): + # Vision-only request (finetune_language_layers=False) with a full leaf list + # must not reach the language-model experts either. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = ["q_proj", "gate_proj", "up_proj", "down_proj"] + scoped_regex = ( + r"(?:.*?(?:vision|visual|image).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = False, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_in_scope_mlp_full_list_still_discovers_moe_parameters(): + # With MLP and language both in scope, an explicit list that names MLP + # leaves SHOULD enable the experts (unchanged behavior): the original list + # is preferred and carries the gate/up/down intent. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + scoped_regex = r".*self_attn.*proj" # unused: original list is preferred + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is original_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_attention_only_list_prefers_original_when_in_scope(): + # The case the PR originally fixed: an attention-only list routed through + # get_peft_regex under a family scope (e.g. vision-off) still keeps experts + # off, because with MLP+language in scope detection uses the original + # attention-only list rather than the regex's spurious mlp component block. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + scoped_regex = ( # carries the spurious mlp block get_peft_regex always adds + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + attn_only_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is attn_only_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None diff --git a/tests/test_prefetch_snapshot_scope.py b/tests/test_prefetch_snapshot_scope.py new file mode 100644 index 0000000000..c7ec4f2c34 --- /dev/null +++ b/tests/test_prefetch_snapshot_scope.py @@ -0,0 +1,916 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""Pure-CPU, no-network unit tests for prefetch snapshot scoping in unsloth/models/_utils.py. + +maybe_prefetch_hf_snapshot warms the HF cache before the in-process load. The warm must cover at +least what the load reads (else the missing file falls to an unprotected in-process Xet fetch) but +not pull weights the load never reads. These tests lock the allow/ignore patterns each mode hands +snapshot_download_with_xet_fallback. The zoo downloader is monkeypatched to capture its kwargs. +""" + +import fnmatch +import sys +import types + +import pytest + +from unsloth.models import _utils as U + + +def _filter(names, allow_patterns, ignore_patterns): + """Mirror HF filter_repo_objects: keep on allow match (or None), drop on ignore match.""" + kept = [] + for name in names: + if allow_patterns is not None and not any(fnmatch.fnmatch(name, p) for p in allow_patterns): + continue + if ignore_patterns and any(fnmatch.fnmatch(name, p) for p in ignore_patterns): + continue + kept.append(name) + return kept + + +@pytest.fixture +def capture(monkeypatch): + """Run maybe_prefetch_hf_snapshot with a fake repo, capturing the patterns forwarded to a + fake injected zoo downloader (independent of the installed unsloth_zoo). Offline env cleared.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + + state = {} + + def fake_download(repo_id, **kw): + state["repo_id"] = repo_id + state["allow_patterns"] = kw.get("allow_patterns") + state["ignore_patterns"] = kw.get("ignore_patterns") + state["variant"] = kw.get("variant") + return "/tmp/fake-snapshot" + + fake_module = types.ModuleType("unsloth_zoo.hf_xet_fallback") + fake_module.snapshot_download_with_xet_fallback = fake_download + fake_module.DownloadStallError = type("DownloadStallError", (RuntimeError,), {}) + monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_fallback", fake_module) + + # Neutralize the model_info network call by default; tests exercising format selection + # install their own. + import huggingface_hub + + class _NoNetworkApi: + def model_info(self, *a, **k): + raise RuntimeError("no network in test") + + monkeypatch.setattr(huggingface_hub, "HfApi", _NoNetworkApi) + + def run(**call_kwargs): + state.clear() + ok = U.maybe_prefetch_hf_snapshot("some-org/some-repo", **call_kwargs) + return ok, state + + return run + + +# Representative repo listing: root weights + aux, subdir, adapter, checkpoint, merged weights. +_SAMPLE_FILES = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + "pytorch_model.bin", + "fp16/model.safetensors", + "experimental/model-00001-of-00002.safetensors", + "checkpoint-500/model.safetensors", + "adapter_config.json", + "adapter_model.safetensors", +] + + +def test_weights_at_root_excludes_subdir_weights(capture): + """A root load ignores subdir weights (fp16/, experimental/, checkpoint-500/) but keeps root weights.""" + ok, st = capture(weights_at_root = True, use_safetensors = True) + assert ok is True + assert st["allow_patterns"] is None + ig = st["ignore_patterns"] + assert "*/*.safetensors" in ig and "*/*.bin" in ig + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], ig) + assert "model-00001-of-00002.safetensors" in kept + assert "model.safetensors.index.json" in kept + assert "config.json" in kept + assert "fp16/model.safetensors" not in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + assert "checkpoint-500/model.safetensors" not in kept + + +def test_adapter_only_excludes_merged_weights(capture): + """An adapter warm keeps adapter files + root aux, not merged full-model weights.""" + ok, st = capture(adapter_only = True) + assert ok is True + assert st["ignore_patterns"] is None + allow = st["allow_patterns"] + assert "adapter_config.json" in allow and "adapter_model*" in allow + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "adapter_config.json" in kept + assert "adapter_model.safetensors" in kept + assert "config.json" in kept and "tokenizer.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "pytorch_model.bin" not in kept + assert "fp16/model.safetensors" not in kept + + +def test_adapter_only_warms_sharded_adapter(capture): + """A sharded adapter is still covered by the adapter_model* glob.""" + _, st = capture(adapter_only = True) + sharded = [ + "adapter_config.json", + "adapter_model-00001-of-00002.safetensors", + "adapter_model-00002-of-00002.safetensors", + "adapter_model.safetensors.index.json", + ] + kept = _filter(sharded, st["allow_patterns"], st["ignore_patterns"]) + assert set(kept) == set(sharded) + + +def test_tokenizer_only_warms_only_aux_files(capture): + """A tokenizer-only repo warms tokenizer/config/vocab files, never weights.""" + _, st = capture(tokenizer_only = True) + assert st["ignore_patterns"] is None + assert st["allow_patterns"] == list(U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "tokenizer.json" in kept and "config.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "adapter_model.safetensors" not in kept + + +def test_aux_warm_covers_arbitrary_remote_code_modules(capture): + """The aux warm must cover any *.py, since trust_remote_code auto_map names modules freely.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + assert "*.py" in allow + remote_code = [ + "config.json", + "modeling.py", + "tokenization.py", + "my_custom_code.py", + "configuration_foo.py", + ] + kept = _filter(remote_code, allow, st["ignore_patterns"]) + for name in ("modeling.py", "tokenization.py", "my_custom_code.py", "configuration_foo.py"): + assert name in kept, name + + +def test_subfolder_warms_subfolder_plus_root_aux(capture): + """A subfolder load warms that subfolder's weights plus root aux; other subdirs/root weights skipped.""" + _, st = capture(subfolder = "fp16") + allow = st["allow_patterns"] + assert "fp16/*" in allow + assert all(p in allow for p in U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + assert "config.json" in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + + +def test_subfolder_takes_precedence_over_weights_at_root(capture): + """When a subfolder is requested the subfolder branch wins over weights_at_root.""" + _, st = capture(subfolder = "fp16", weights_at_root = True) + assert "fp16/*" in st["allow_patterns"] + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + + +def test_local_dir_is_not_warmed(capture, tmp_path): + """A local directory path skips the warm (returns False).""" + d = tmp_path / "local-model" + d.mkdir() + ok = U.maybe_prefetch_hf_snapshot(str(d), weights_at_root = True) + assert ok is False + + +def _install_fake_model_info(monkeypatch, filenames): + """Make HfApi().model_info(...).siblings report filenames, with no network.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Info: + def __init__(self, names): + self.siblings = [_Sib(n) for n in names] + + class _Api: + def model_info(self, *a, **k): + return _Info(filenames) + + monkeypatch.setattr(huggingface_hub, "HfApi", _Api) + + +# ----- Finding P: variant-aware weight-format selection ----- + + +def test_variant_keeps_bin_when_only_default_safetensors(monkeypatch): + """A default model.safetensors must not prove a variant .bin redundant; without a variant it does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + ig_default = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_default + + +def test_variant_drops_bin_when_variant_safetensors_present(monkeypatch): + """A variant-matching safetensors makes the variant .bin redundant, so .bin is dropped.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_no_variant_keeps_bin_when_only_variant_safetensors(monkeypatch): + """For a no-variant load, only a canonical safetensors (not a lone variant) makes .bin redundant.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_variant_keeps_bin_for_noncanonical_sidecar(monkeypatch): + """A non-canonical variant sidecar must not prove the variant .bin redundant; a canonical one does.""" + _install_fake_model_info( + monkeypatch, ["consolidated.fp16.safetensors", "pytorch_model.fp16.bin"] + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_is_canonical_model_weight_safetensors(): + """The canonical detector matches only non-variant model-weight safetensors names.""" + assert U._is_canonical_model_weight_safetensors("model.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_canonical_model_weight_safetensors("model.fp16.safetensors") is False + assert ( + U._is_canonical_model_weight_safetensors("model.fp16-00001-of-00002.safetensors") is False + ) + assert U._is_canonical_model_weight_safetensors("adapter_model.safetensors") is False + + +def test_st_prefetch_resolves_env_cache_and_runs_after_validation(): + """The ST prefetch must resolve SENTENCE_TRANSFORMERS_HOME and run after load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert len(prefetch_calls) == 1, "expected exactly one ST prefetch call" + call = prefetch_calls[0] + # cache_dir kwarg resolves SENTENCE_TRANSFORMERS_HOME. + cache_dir_kw = next((kw for kw in call.keywords if kw.arg == "cache_dir"), None) + assert cache_dir_kw is not None, "ST prefetch must pass cache_dir" + assert "SENTENCE_TRANSFORMERS_HOME" in ast.dump( + cache_dir_kw.value + ), "ST prefetch cache_dir must resolve SENTENCE_TRANSFORMERS_HOME" + # Load-mode validation runs before the prefetch (fewer source lines = earlier). + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < call.lineno, "load-mode validation must precede the ST prefetch" + + +def test_st_cache_resolutions_honor_explicit_hf_cache_dir(): + """Every ST cache resolution falling back to SENTENCE_TRANSFORMERS_HOME must first honor an explicit HF cache_dir.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + resolutions = [ + kw + for kw in ast.walk(tree) + if isinstance(kw, ast.keyword) + and kw.arg == "cache_dir" + and "SENTENCE_TRANSFORMERS_HOME" in ast.dump(kw.value) + ] + assert resolutions, "expected cache_dir resolutions referencing SENTENCE_TRANSFORMERS_HOME" + for kw in resolutions: + assert "'cache_dir'" in ast.dump( + kw.value + ), "an ST cache_dir resolution must read an explicit kwargs.get('cache_dir') first" + + +def test_st_native_loads_map_hf_cache_dir_to_cache_folder(): + """Native SentenceTransformer loads take cache_folder, so an explicit HF cache_dir must be mapped onto it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + # Every native SentenceTransformer(...) forwarding cache_folder must read cache_dir. + st_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ] + cache_folder_kws = [kw for call in st_calls for kw in call.keywords if kw.arg == "cache_folder"] + assert cache_folder_kws, "expected a native SentenceTransformer call forwarding cache_folder" + for kw in cache_folder_kws: + assert "'cache_dir'" in ast.dump( + kw.value + ), "a native SentenceTransformer cache_folder must map the explicit HF cache_dir first" + # for_inference feeds cache_folder via st_kwargs; both native branches map cache_dir -> cache_folder. + normalized = "".join(src.split()) + assert ( + 'st_kwargs["cache_folder"]=' in normalized + ), "for_inference must set st_kwargs cache_folder" + assert ( + normalized.count('kwargs.get("cache_dir")orkwargs.get("cache_folder")') >= 2 + ), "both native ST branches (for_inference, fast-encoder) must map cache_dir -> cache_folder" + + +def test_vision_warms_vllm_tokenizer_after_remap(): + """On the vLLM path the tokenizer warm is deferred until after the fast_inference_setup remap.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + guard = "if _vllm_owns_weights and isinstance(tokenizer_name" + assert guard in src, "expected a vLLM-gated tokenizer warm" + assert src.index(guard) > src.index( + "fast_inference_setup(" + ), "the vLLM tokenizer warm must run after the fast_inference_setup remap" + + +def test_diffusion_forwards_variant_to_real_load(): + """FastDiffusionModel must forward variant to the real model_cls.from_pretrained load, not just the prefetch.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "diffusion.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + assert ( + 'load_kwargs["variant"] = kwargs["variant"]' in src + ), "the diffusion load must forward variant to model_cls.from_pretrained" + + +def test_vision_prefetch_runs_after_load_mode_validation(): + """The FastBaseModel (vision) prefetch must run after the load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert prefetch_calls, "expected a vision prefetch call" + first_prefetch = min(call.lineno for call in prefetch_calls) + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < first_prefetch, "load-mode validation must precede the vision prefetch" + + +def test_llama_prefetch_skips_only_real_vllm_loads(): + """The llama prefetch's fast_inference skip must be gated on num_labels is None (a classification load still downloads).""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "llama.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + gated = False + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ): + continue + fi_kw = next((kw for kw in n.keywords if kw.arg == "fast_inference"), None) + if fi_kw is None: + continue + dumped = ast.dump(fi_kw.value) + if "fast_inference" in dumped and "num_labels" in dumped: + gated = True + assert gated, "llama prefetch fast_inference must be gated on num_labels is None" + + +def test_st_fallback_module_loads_resolve_env_cache(): + """Fallback module loads deriving cache_dir from cache_folder must also fall back to SENTENCE_TRANSFORMERS_HOME.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + + # Fallback sites (cache_dir derived from cache_folder) must resolve SENTENCE_TRANSFORMERS_HOME. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None: + continue + dumped = ast.dump(cache_dir_kw.value) + if "cache_folder" not in dumped: + continue # internal pass-through, not a resolution site + checked += 1 + assert ( + "SENTENCE_TRANSFORMERS_HOME" in dumped + ), f"{node.func.attr} cache_dir resolves cache_folder but not SENTENCE_TRANSFORMERS_HOME" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to resolve the env cache" + + +def test_st_fallback_module_loads_forward_revision(): + """The fallback module loads must forward revision so module files match the revision-pinned weights. + Guards: (a) helpers accept revision, (b) every download primitive forwards it, (c) _load_modules + threads it into internal calls, (d) the from_pretrained fallback sites forward it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + funcs = { + n.name: n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) + and n.name in ("_module_path", "_read_pooling_mode", "_load_modules") + } + assert set(funcs) == {"_module_path", "_read_pooling_mode", "_load_modules"} + + # (a) each helper takes a revision parameter. + for name, fn in funcs.items(): + arg_names = {a.arg for a in fn.args.args + fn.args.kwonlyargs} + assert "revision" in arg_names, f"{name} must accept a revision argument" + + # (b) every download primitive inside the helpers forwards revision. + downloads = 0 + for name, fn in funcs.items(): + for node in ast.walk(fn): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)): + continue + if node.func.id not in ("hf_hub_download", "load_dir_path"): + continue + downloads += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"{node.func.id} in {name} must forward revision" + assert downloads >= 3, "expected the module-download primitives to be revision-guarded" + + # (c) _load_modules threads revision into its internal _module_path / _read_pooling_mode calls. + internal = 0 + for node in ast.walk(funcs["_load_modules"]): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_read_pooling_mode"): + continue + internal += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"_load_modules must forward revision to {node.func.attr}" + assert internal >= 2, "expected _load_modules to call _module_path and _read_pooling_mode" + + # (d) the from_pretrained fallback _module_path / _load_modules sites forward revision. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None or "cache_folder" not in ast.dump(cache_dir_kw.value): + continue # internal pass-through, not a fallback site + checked += 1 + rev_kw = next((kw for kw in node.keywords if kw.arg == "revision"), None) + assert rev_kw is not None and "revision" in ast.dump( + rev_kw.value + ), f"{node.func.attr} fallback call must forward revision" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to forward revision" + + +def test_st_fallback_model_load_resolves_env_cache(): + """from_pretrained must resolve the warmed ST cache into kwargs['cache_dir'] before the FastModel weight load.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + def _resolves_st_cache(value_node): + # Resolution may be inline or in the assignment to an intermediate variable the value references. + dumped = ast.dump(value_node) + if "cache_folder" in dumped and "SENTENCE_TRANSFORMERS_HOME" in dumped: + return True + if isinstance(value_node, ast.Name): + for n in ast.walk(tree): + if isinstance(n, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == value_node.id for t in n.targets + ): + d = ast.dump(n.value) + if "cache_folder" in d and "SENTENCE_TRANSFORMERS_HOME" in d: + return True + return False + + resolved_lines = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + for tgt in node.targets: + if ( + isinstance(tgt, ast.Subscript) + and isinstance(tgt.value, ast.Name) + and tgt.value.id == "kwargs" + and isinstance(tgt.slice, ast.Constant) + and tgt.slice.value == "cache_dir" + and _resolves_st_cache(node.value) + ): + resolved_lines.append(node.lineno) + assert resolved_lines, "from_pretrained must resolve the ST cache into kwargs['cache_dir']" + + fastmodel_calls = [ + n.lineno + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "from_pretrained" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "FastModel" + ] + assert fastmodel_calls, "expected a FastModel.from_pretrained call" + assert min(resolved_lines) < min( + fastmodel_calls + ), "kwargs['cache_dir'] must be resolved before the fallback FastModel weight load" + + +def test_canonical_variant_model_weight_matches_transformers_names(): + """The variant safetensors detector matches only canonical variant names, rejecting sidecars and wrong variants.""" + f = U._is_canonical_variant_model_weight_safetensors + assert f("model.fp16.safetensors", "fp16") is True + assert f("model.fp16-00001-of-00002.safetensors", "fp16") is True + assert f("model-00001-of-00002.fp16.safetensors", "fp16") is True + assert f("model.safetensors.index.fp16.json", "fp16") is True + assert f("consolidated.fp16.safetensors", "fp16") is False + assert f("model.safetensors", "fp16") is False + assert f("model-00001-of-00002.safetensors", "fp16") is False + assert f("model.bf16.safetensors", "fp16") is False + + +def test_variant_is_forwarded_to_downloader(capture): + """maybe_prefetch_hf_snapshot must forward variant to the downloader (absent a variant, nothing is forwarded).""" + _, st = capture(weights_at_root = True, use_safetensors = True, variant = "fp16") + assert st["variant"] == "fp16" + _, st = capture(weights_at_root = True, use_safetensors = True) + assert st["variant"] is None + + +def test_variant_drops_bin_for_sharded_variant_safetensors(monkeypatch): + """A sharded variant safetensors is recognized, so its redundant variant .bin is dropped.""" + _install_fake_model_info( + monkeypatch, + [ + "model.fp16-00001-of-00002.safetensors", + "model.fp16-00002-of-00002.safetensors", + "pytorch_model.fp16-00001-of-00002.bin", + ], + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_tokenizer_only_warms_extra_vocab_files(capture): + """tokenizer_only must warm SentencePiece / vocab / processor files, including a named jinja template.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + ): + assert name in allow, name + sample = [ + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + "additional_chat_templates/custom.jinja", + ] + kept = _filter(sample, allow, st["ignore_patterns"]) + assert set(kept) == set(sample) + + +def test_format_probe_runs_even_when_config_cached(capture, monkeypatch): + """A cached config.json must not skip the weight-format probe; model_info still drops the redundant .bin.""" + import huggingface_hub + + # Pretend config.json is cached (the AutoConfig side effect); this must not gate the probe. + monkeypatch.setattr( + huggingface_hub, "try_to_load_from_cache", lambda *a, **k: "/cache/config.json" + ) + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + _, st = capture(weights_at_root = True) + ig = st["ignore_patterns"] or [] + assert "*.bin" in ig + + +def test_optimizer_safetensors_does_not_drop_bin(monkeypatch): + """An optimizer.safetensors sidecar must not count as model safetensors, so the real .bin weights are kept.""" + _install_fake_model_info(monkeypatch, ["pytorch_model.bin", "optimizer.safetensors"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + + +def test_model_safetensors_still_drops_bin(monkeypatch): + """Control for the optimizer case: a real model.safetensors next to pytorch_model.bin still drops the .bin.""" + _install_fake_model_info( + monkeypatch, ["model.safetensors", "pytorch_model.bin", "optimizer.safetensors"] + ) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig + + +def test_whole_multi_component_snapshot_keeps_subdir_bin(monkeypatch): + """A whole multi-component snapshot must not drop *.bin (it would strip a subdir module's weight); a root load still does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "1_Dense/pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = False) + assert "*.bin" not in ig + ig_root = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_root + + +def test_is_model_weight_safetensors_classification(): + """Real model weights count; adapter / trainer-state sidecars do not.""" + assert U._is_model_weight_safetensors("model.safetensors") is True + assert U._is_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_model_weight_safetensors("consolidated.safetensors") is True + assert U._is_model_weight_safetensors("adapter_model.safetensors") is False + assert U._is_model_weight_safetensors("optimizer.safetensors") is False + assert U._is_model_weight_safetensors("scheduler.safetensors") is False + assert U._is_model_weight_safetensors("rng_state_0.safetensors") is False + + +def test_tokenizer_only_warms_slow_sentencepiece_vocab(capture): + """tokenizer_only must warm the slow-tokenizer SentencePiece / BPE vocab files AutoTokenizer fetches first.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + ): + assert name in allow, name + + +def test_adapter_safetensors_check_scoped_to_root(monkeypatch): + """_adapter_repo_has_safetensors must only count a root adapter_model*.safetensors, not a subdir one.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Api: + def __init__(self, names): + self._names = names + + def model_info(self, *a, **k): + return type("MI", (), {"siblings": [_Sib(n) for n in self._names]})() + + # Subdir safetensors only -> not reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api( + ["adapter_config.json", "adapter_model.bin", "checkpoint-5/adapter_model.safetensors"] + ), + ) + assert U._adapter_repo_has_safetensors("org/repo") is False + # Root safetensors -> reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api(["adapter_config.json", "adapter_model.safetensors"]), + ) + assert U._adapter_repo_has_safetensors("org/repo") is True + + +def test_gguf_file_warm_keeps_gguf(capture): + """A gguf_file load allow-lists that GGUF while not pulling other quants the repo publishes.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf") + allow = st["allow_patterns"] + ig = st["ignore_patterns"] + assert allow is not None and "model-Q4_K_M.gguf" in allow + sample = [ + "model-Q4_K_M.gguf", + "model-Q8_0.gguf", + "config.json", + "tokenizer.json", + ] + kept = _filter(sample, allow, ig) + assert "model-Q4_K_M.gguf" in kept + assert "config.json" in kept + assert "model-Q8_0.gguf" not in kept + + +# ----- Finding Q: adapter weight-format selection ----- + + +def test_adapter_only_prefers_safetensors_over_bin(capture, monkeypatch): + """A mixed-format adapter repo warms only the safetensors PeftModel reads, not both formats.""" + _install_fake_model_info( + monkeypatch, ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"] + ) + _, st = capture(adapter_only = True) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.bin" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.safetensors" in kept + assert "adapter_model.bin" not in kept + + +def test_adapter_only_bin_only_keeps_bin(capture, monkeypatch): + """A .bin-only adapter repo must keep adapter_model.bin (no safetensors found -> both formats eligible).""" + _install_fake_model_info(monkeypatch, ["adapter_config.json", "adapter_model.bin"]) + _, st = capture(adapter_only = True) + kept = _filter( + ["adapter_config.json", "adapter_model.bin"], st["allow_patterns"], st["ignore_patterns"] + ) + assert "adapter_model.bin" in kept + + +def test_adapter_only_explicit_use_safetensors_false_keeps_bin(capture): + """An explicit use_safetensors=False forces the .bin form without a model_info call.""" + _, st = capture(adapter_only = True, use_safetensors = False) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.safetensors" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.bin" in kept + assert "adapter_model.safetensors" not in kept + + +def test_gguf_file_with_subfolder_warms_subfolder_path(capture): + """gguf_file + subfolder: the warm allow-lists /, not the bare root name.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf", subfolder = "gguf") + allow = st["allow_patterns"] + assert "gguf/model-Q4_K_M.gguf" in allow + kept = _filter(["gguf/model-Q4_K_M.gguf", "config.json"], allow, st["ignore_patterns"]) + assert "gguf/model-Q4_K_M.gguf" in kept and "config.json" in kept + + +def test_from_tf_root_load_ignores_nested_h5(capture): + """A from_tf root load keeps the root .h5 but drops nested .h5 / .msgpack checkpoints.""" + _, st = capture(weights_at_root = True, from_tf = True) + ig = st["ignore_patterns"] + assert "*/*.h5" in ig and "*/*.msgpack" in ig + kept = _filter(["model.h5", "checkpoint-1/model.h5", "config.json"], st["allow_patterns"], ig) + assert "model.h5" in kept + assert "checkpoint-1/model.h5" not in kept + + +def test_sentence_transformer_from_pretrained_is_prefetch_wired(): + """from_pretrained must call maybe_prefetch_hf_snapshot as an unconditional top-level statement before any return.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + cls = next( + n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "FastSentenceTransformer" + ) + fp = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "from_pretrained") + + def _prefetch_call(node): + # a bare call statement, or one whose return is captured (e.g. _st_prefetched = ...) + value = node.value if isinstance(node, (ast.Expr, ast.Assign)) else None + if ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id == "maybe_prefetch_hf_snapshot" + ): + return value + return None + + prefetch_pos = next((i for i, n in enumerate(fp.body) if _prefetch_call(n)), None) + return_pos = next((i for i, n in enumerate(fp.body) if isinstance(n, ast.Return)), len(fp.body)) + assert ( + prefetch_pos is not None + ), "from_pretrained must call maybe_prefetch_hf_snapshot at top level" + assert prefetch_pos < return_pos, "prefetch must run before any top-level return" + # local_files_only must be forwarded so an offline load does not start a Hub download. + prefetch_call = _prefetch_call(fp.body[prefetch_pos]) + assert "local_files_only" in { + kw.arg for kw in prefetch_call.keywords + }, "prefetch must forward local_files_only" + + +def test_st_module_download_forwards_cache_folder(): + """_load_modules must forward the custom cache_folder into load_dir_path so per-module subdirs read the warmed cache.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "load_dir_path" + ] + assert calls, "expected a load_dir_path call in sentence_transformer.py" + assert all( + "cache_folder" in {kw.arg for kw in c.keywords} for c in calls + ), "every load_dir_path call must forward cache_folder" + + +def test_st_native_sentence_transformer_calls_forward_cache_folder(): + """Every native SentenceTransformer(model_name, ...) load must forward cache_folder; a modules-based build needs none.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + weight_loading_calls = [] + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ): + continue + kw_names = {kw.arg for kw in n.keywords} + # A modules-based build downloads nothing; only a repo-name load reads the cache. + if "modules" in kw_names: + continue + weight_loading_calls.append(n) + assert ( + weight_loading_calls + ), "expected a repo-name SentenceTransformer load in sentence_transformer.py" + # cache_folder is forwarded explicitly or via a **kwargs unpacking (kw.arg == None). + for c in weight_loading_calls: + kw_names = {kw.arg for kw in c.keywords} + forwards = "cache_folder" in kw_names or None in kw_names + assert forwards, ( + "a repo-name SentenceTransformer load must forward cache_folder " + f"(explicitly or via **kwargs) at line {c.lineno}" + ) diff --git a/tests/test_synthetic_chunk_data.py b/tests/test_synthetic_chunk_data.py index b9167d214f..abc2c01443 100644 --- a/tests/test_synthetic_chunk_data.py +++ b/tests/test_synthetic_chunk_data.py @@ -104,10 +104,36 @@ def test_chunk_data_rejects_overlap_not_smaller_than_chunk(): os.unlink(path) +def test_chunk_data_uninitialized_error_names_real_class(): + # Without max_seq_length the guard tells the user which method to call first. + # The message must name the real class (SyntheticDataKit) so copying it works; + # a misspelling would raise NameError when the user follows it verbatim. + kit = SyntheticDataKit.__new__(SyntheticDataKit) + kit.tokenizer = _MockTokenizer() # max_seq_length intentionally unset + with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f: + f.write("word " * 50) + path = f.name + try: + try: + kit.chunk_data(filename = path) + raise AssertionError("expected RuntimeError when max_seq_length is unset") + except RuntimeError as e: + msg = str(e) + assert ( + "SyntheticDataKit.from_pretrained" in msg + ), f"error must name SyntheticDataKit.from_pretrained, got: {msg}" + assert ( + "SynthetidDataKit" not in msg + ), f"error must not misspell the class name, got: {msg}" + finally: + os.unlink(path) + + if __name__ == "__main__": test_chunk_data_keeps_single_chunk_document() test_chunk_data_still_splits_long_document() test_chunk_data_empty_document_yields_no_chunks() test_chunk_data_short_document_is_not_split_into_fragments() test_chunk_data_rejects_overlap_not_smaller_than_chunk() + test_chunk_data_uninitialized_error_names_real_class() print("OK") diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index b976654f87..98f7e2db62 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -31,6 +31,7 @@ requires_cuda = pytest.mark.skipif( REPO_ROOT = Path(__file__).resolve().parents[2] LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py" +LOADER_PY = REPO_ROOT / "unsloth" / "models" / "loader.py" CLASS_NAME = "LlamaRotaryEmbedding" @@ -78,42 +79,88 @@ def _config_branch(init_fn): return None +def _iter_names_and_calls(node): + """(attribute/string names, bare-name calls, method-call attrs) under node.""" + names, calls, call_attrs = set(), set(), set() + for sub in ast.walk(node): + if isinstance(sub, ast.Attribute): + names.add(sub.attr) + elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): + names.add(sub.value) + elif isinstance(sub, ast.Call): + if isinstance(sub.func, ast.Name): + calls.add(sub.func.id) + elif isinstance(sub.func, ast.Attribute): + call_attrs.add(sub.func.attr) + return names, calls, call_attrs + + +def _find_method(source_path, class_name, method_name): + for node in ast.walk(ast.parse(source_path.read_text())): + if isinstance(node, ast.ClassDef) and node.name == class_name: + for sub in node.body: + if isinstance(sub, ast.FunctionDef) and sub.name == method_name: + return sub + return None + + +def _find_function(source_path, function_name): + for node in ast.walk(ast.parse(source_path.read_text())): + if isinstance(node, ast.FunctionDef) and node.name == function_name: + return node + return None + + def test_config_path_inspects_rope_scaling(): init_fn = _load_class_init() - branch = _config_branch(init_fn) - assert branch is not None, ( - f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` " - "branch; the config constructor path must read config.rope_scaling so " - "scaled models (llama3/linear/longrope) are not silently unscaled " - "(issue #2405)" - ) + # inv_freq is derived through the shared _unsloth_recompute_inv_freq helper + # (or still inlined in the config branch on older layouts); whichever scope + # holds the scaling must read config.rope_scaling and call + # _compute_config_rope_inv_freq, else scaled models run unscaled (#2405). + _, _, init_call_attrs = _iter_names_and_calls(init_fn) + scope = _find_method(LLAMA_PY, CLASS_NAME, "_unsloth_recompute_inv_freq") + if scope is not None: + assert "_unsloth_recompute_inv_freq" in init_call_attrs, ( + f"{CLASS_NAME}.__init__ no longer derives inv_freq via " + "_unsloth_recompute_inv_freq; keep the constructor wired to the " + "shared scaling helper or scaled configs silently lose RoPE scaling " + "(issue #2405)." + ) + else: + scope = _config_branch(init_fn) + assert scope is not None, ( + f"{CLASS_NAME}.__init__ has neither a _unsloth_recompute_inv_freq " + "helper nor an `if config is not None:` branch; the config path must " + "apply llama3/linear/longrope scaling (issue #2405)." + ) - names = set() - for stmt in branch.body: - for sub in ast.walk(stmt): - if isinstance(sub, ast.Attribute): - names.add(sub.attr) - elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): - names.add(sub.value) + names, called, _ = _iter_names_and_calls(scope) assert "rope_scaling" in names, ( - f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. " - "When a rotary class is built straight from a config (the path modern " - "transformers takes, since rotary moved to LlamaModel), the llama3 / " - "linear / longrope scaling must still be applied; otherwise long inputs " - "produce repeated-pattern gibberish (issue #2405)." + f"{CLASS_NAME} inv_freq computation does not reference `rope_scaling`; " + "scaled models (llama3/linear/longrope) would run unscaled and produce " + "repeated-pattern gibberish past the original context (issue #2405)." + ) + assert "_compute_config_rope_inv_freq" in called, ( + f"{CLASS_NAME} inv_freq computation no longer calls " + "_compute_config_rope_inv_freq; keep it wired or scaled configs silently " + "lose RoPE scaling again (issue #2405)." ) - called = { - sub.func.id - for stmt in branch.body - for sub in ast.walk(stmt) - if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name) - } - assert "_compute_config_rope_inv_freq" in called, ( - f"{CLASS_NAME}.__init__ config path no longer calls " - "_compute_config_rope_inv_freq; the CPU behavioral tests below cover " - "that helper directly, so the constructor must stay wired to it or " - "scaled configs silently lose RoPE scaling again (issue #2405)." + +def test_v5_repair_reuses_recompute(): + # transformers v5 blanks non-persistent buffers on load, so + # loader._fix_rope_inv_freq rebuilds inv_freq; it must reuse the scaled + # recompute, since an unscaled rebuild re-drops llama3 scaling (#2405). + fix_fn = _find_function(LOADER_PY, "_fix_rope_inv_freq") + assert fix_fn is not None, ( + "loader._fix_rope_inv_freq not found; if it was renamed, update this " + "guard so the v5 rope repair keeps applying config scaling (issue #2405)." + ) + _, _, call_attrs = _iter_names_and_calls(fix_fn) + assert "_unsloth_recompute_inv_freq" in call_attrs, ( + "loader._fix_rope_inv_freq no longer rebuilds inv_freq via " + "_unsloth_recompute_inv_freq; transformers v5 blanks the buffer on load " + "and an unscaled rebuild re-drops llama3 scaling (issue #2405)." ) @@ -189,6 +236,27 @@ def test_default_rope_type_matches_vanilla_inv_freq(): ) +def test_recompute_helper_scales_on_cpu(): + # Exercise the exact method loader._fix_rope_inv_freq calls, without CUDA. + from unsloth.models.llama import LlamaRotaryEmbedding, _get_rope_theta + + def recompute(config): + rot = object.__new__(LlamaRotaryEmbedding) + rot.attention_scaling = 1.0 + rot.base = _get_rope_theta(config, 10000.0) + rot.dim = config.head_dim + rot._unsloth_rope_config = config + return rot._unsloth_recompute_inv_freq().float().cpu() + + config = _make_config(LLAMA3_ROPE_SCALING) + assert torch.allclose( + recompute(config), _reference_inv_freq(config, "llama3"), rtol = 1e-4, atol = 1e-6 + ), "_unsloth_recompute_inv_freq dropped llama3 scaling (issue #2405)." + assert torch.allclose( + recompute(_make_config(None)), _vanilla_inv_freq(), rtol = 1e-4, atol = 1e-6 + ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." + + def _cos_at_position(rot, position): """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" inv_freq = rot.inv_freq.float().cpu() diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index f0a843c380..8f7ed6e09e 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -203,6 +203,7 @@ def main(): ap.add_argument("--max-seq-length", type = int, default = 2048) ap.add_argument("--is-vlm", action = "store_true") ap.add_argument("--trust-remote-code", action = "store_true") + ap.add_argument("--trust-remote-code-tokenizer", action = "store_true") ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards") args = ap.parse_args() @@ -232,7 +233,11 @@ def main(): model.eval() # A tokenizer may be absent if the caller saved it separately; only calibration needs one. try: - tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code) + # The tokenizer/processor has its own trust flag: consent for one component must not + # let the other's custom code run. + tokenizer = auto_proc.from_pretrained( + args.model, trust_remote_code = args.trust_remote_code_tokenizer + ) except Exception: if args.needs_calibration: raise RuntimeError( diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 10690810fe..6f025343f5 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -391,7 +391,7 @@ class SyntheticDataKit: assert os.path.exists(filename) assert hasattr(self, "tokenizer") if not hasattr(self, "max_seq_length"): - raise RuntimeError("Please use SynthetidDataKit.from_pretrained(...) first!") + raise RuntimeError("Please use SyntheticDataKit.from_pretrained(...) first!") if not hasattr(self, "overlap") or not hasattr(self, "max_generation_tokens"): raise RuntimeError("Please use prepare_qa_generation first!") diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index ca608fa01b..935ffbb447 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -68,7 +68,9 @@ def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): n = tl.cdiv(N, BLOCK_SIZE) offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - offs = offs_m[:, None] * N + offs_n[None, :] + # tl.arange is int32, so offs_m * N overflows for tensors with more than + # 2**31 elements (e.g. flattened MoE expert stacks); index in int64. + offs = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64) mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) x = tl.load(x_ptr + offs, mask = mask).to(tl.float32) s = tl.load(s_ptr + pid_m * n + pid_n) @@ -327,11 +329,42 @@ fp8_block_matmul = ( ) +def _blockwise_weight_dequant_any_shape(weight, weight_scale, block_size, out_dtype): + """Blockwise fp8 weight dequant for any shape: triton when the weight tiles + evenly into block_size, else a torch-native per-block scale expansion.""" + m, n = weight.shape + if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16): + weight_scale = weight_scale.to(torch.float32) # e.g. float8_e8m0fnu scales break triton + if weight_scale.numel() == 1: + # Per-tensor scale: the normal forward stashes the un-expanded scalar, + # which repeat_interleave cannot grow to (m, n). Scale directly. + return (weight.to(torch.float32) * weight_scale.float()).to(out_dtype) + if m % block_size[0] != 0 or n % block_size[1] != 0 or block_size[0] != block_size[1]: + # Uneven tiling, or rectangular blocks. The triton kernel uses a single + # BLOCK_SIZE for both axes and derives the column scale stride from it, so + # it mis-indexes the scale when block_size[0] != block_size[1]. Expand the + # per-block scales in torch, which handles both dimensions independently. + s_full = weight_scale.repeat_interleave(block_size[0], 0)[:m] + s_full = s_full.repeat_interleave(block_size[1], 1)[:, :n] + return (weight.to(torch.float32) * s_full).to(out_dtype) + # Even tiling with square blocks: block-quant dequant with the real block size + # (weight_dequant would silently default to 128 and dequantize wrongly). + return weight_dequant_block(weight, weight_scale, block_size = block_size[0], dtype = out_dtype) + + class FP8BlockQuantLinear(torch.autograd.Function): @staticmethod def forward(ctx, X, weight, weight_scale): m, n = weight.shape + if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16): + # Upcast (e.g. e8m0) returns a fresh tensor and drops any Python + # attribute, so carry block_size across the cast for the lookup below. + _scale_block_size = getattr(weight_scale, "block_size", None) + weight_scale = weight_scale.to(torch.float32) # e8m0 scales break triton dtype mapping + if _scale_block_size is not None: + weight_scale.block_size = _scale_block_size + # Original scale, saved for backward before any transformation original_weight_scale = weight_scale @@ -360,6 +393,18 @@ class FP8BlockQuantLinear(torch.autograd.Function): if not weight.is_contiguous(): weight = weight.contiguous() + if X.shape[-1] % block_size[1] != 0: + # Hidden dim not divisible by the activation block: dequant + plain matmul. + # Use the original (un-expanded) scale so a scalar per-tensor scale keeps + # the fast scalar path in both forward and backward. + W_deq = _blockwise_weight_dequant_any_shape( + weight, original_weight_scale, block_size, X.dtype + ) + ctx.weight = weight + ctx.weight_scale = original_weight_scale + ctx.block_size = block_size + return torch_matmul(X, W_deq.T).to(X.dtype) + qinput, scale = act_quant(X, block_size[1]) output = fp8_block_matmul( qinput, @@ -371,11 +416,14 @@ class FP8BlockQuantLinear(torch.autograd.Function): ) ctx.weight = weight ctx.weight_scale = original_weight_scale # Save original for backward + ctx.block_size = block_size return output.to(X.dtype) @staticmethod def backward(ctx, grad_output): - W_deq = weight_dequant(ctx.weight, ctx.weight_scale) + W_deq = _blockwise_weight_dequant_any_shape( + ctx.weight, ctx.weight_scale, ctx.block_size, grad_output.dtype + ) grad_X = torch_matmul(grad_output, W_deq) del W_deq return grad_X, None, None diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 047783c35e..169b610988 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -83,8 +83,10 @@ __all__ = [ "verify_fp8_support_if_applicable", "_get_inference_mode_context_manager", "hf_login", + "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", + "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", "_patch_transformers_trainer_data_parallel", @@ -421,6 +423,18 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") _SDPA_EXCLUDED_MODELS = ("gpt_oss",) +# The loader (loader.py) forces supports_sdpa=False for these because their bundled +# SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor +# them without a loader -> _utils import cycle (loader.py already imports from _utils +# and re-exports this name for callers like sentence_transformer.py). Entries are matched +# as substrings against a comma-joined model_types string ending in a comma, so "gemma3," +# matches a distinct "gemma3" entry but not "gemma3n", and "gemma3_text" matches the +# EmbeddingGemma text model. +DISABLE_SDPA_MODEL_NAMES = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore + "gpt_oss", +] _FLASH_EXCLUDED_MODELS = ("gpt_oss",) _EAGER_ONLY_PREFIXES = ("gemma3n",) _FLASH_ATTENTION_MAX_HEAD_DIM = 256 @@ -431,8 +445,23 @@ def _is_flex_excluded(model_type): return model_type in _FLEX_EXCLUDED_MODELS +def _is_sdpa_disabled_by_name(model_type): + # Mirror the loader's DISABLE_SDPA_MODEL_NAMES check: loader.py builds + # model_types_all = ",".join(model_types) + "," and tests `name in model_types_all`. + # Rebuild the same trailing-comma form for a single model_type so the match is + # identical (e.g. "gemma3," matches "gemma3" but not "gemma3n", and "gemma3_text" + # still matches "gemma3_text"). + model_types_all = model_type.lower() + "," + return any(name.lower() in model_types_all for name in DISABLE_SDPA_MODEL_NAMES) + + def _is_sdpa_excluded(model_type): - return model_type in _SDPA_EXCLUDED_MODELS + # SDPA is known-broken for these models, so an explicit sdpa request must not + # re-enable it. Two sources: _SDPA_EXCLUDED_MODELS (resolver-level, e.g. gpt_oss) + # and DISABLE_SDPA_MODEL_NAMES (loader-level, e.g. gemma3 / gemma3_text, which the + # loader also forces to supports_sdpa=False). + lowered = model_type.lower() + return lowered in _SDPA_EXCLUDED_MODELS or _is_sdpa_disabled_by_name(lowered) def _is_flash_excluded(model_type): @@ -608,6 +637,12 @@ def _disable_flash_attention_if_needed( if disable_reason is None: return attn_implementation + # Only an implementation passed by the caller counts as an explicit request. + # Values read from the config are synthesized by the loaders (the language path + # seeds the config with attn_implementation="sdpa") or come from Transformers + # defaults, so they must not be treated as a deliberate user choice. + explicit_request = attn_implementation + requested_attn_implementation = attn_implementation if requested_attn_implementation is None: requested_attn_implementation = _config_get(config, "_attn_implementation", None) @@ -617,6 +652,20 @@ def _disable_flash_attention_if_needed( if requested_attn_implementation == "eager": return _set_attn_impl(config, "eager") + model_type = _config_get(config, "model_type", "") + + # The disable reason is flash-specific: honor an explicit non-flash request from + # the caller instead of downgrading it. SDPA is honored unless the model's SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text); flex_attention + # is honored only when it is actually usable, since supports_flex_attention already + # rejects the excluded/broken/unavailable configs. This keeps an explicit request + # from selecting a backend the repo marks as wrong. + if explicit_request == "sdpa" and not _is_sdpa_excluded(model_type.lower()): + return _set_attn_impl(config, "sdpa") + if explicit_request == "flex_attention" and supports_flex_attention: + return _set_attn_impl(config, "flex_attention") + if supports_sdpa: fallback_attn_implementation = "sdpa" elif supports_flex_attention: @@ -629,7 +678,6 @@ def _disable_flash_attention_if_needed( if _is_flash_attention_requested(requested_attn_implementation) else "flash_attention_2" ) - model_type = _config_get(config, "model_type", "") warning_key = ( model_type, logged_attn_implementation, @@ -843,7 +891,19 @@ def resolve_attention_implementation( final_attn_impl = requested_attn_implementation _set_attn_impl(config, final_attn_impl) - if not supports_sdpa and final_attn_impl == "sdpa": + # A caller who explicitly passes requested_attn_implementation="sdpa" keeps it even + # on a conservatively unsupported model, mirroring _disable_flash_attention_if_needed + # which honors an explicit sdpa request. The exception is a model whose SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text, which the loader also forces to supports_sdpa=False): + # an explicit request must not re-enable it, so it still downgrades to eager, just + # like flex falls back for _FLEX_EXCLUDED_MODELS. A synthesized/default sdpa + # (requested is None, so the value came from the model resolution above or the + # config) also downgrades. + honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded( + model_type + ) + if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa: print( f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager." ) @@ -905,6 +965,411 @@ logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITI TORCHAO_MSG = "Error: torchao not found, please install with `pip install torchao`" +# Artifacts a Transformers/PEFT load never reads (ONNX/TF/Flax/CoreML/GGUF/training state), skipped +# when prewarming so a mixed-format repo is not pulled in full. +_PREFETCH_IGNORE_PATTERNS = ( + "*.onnx", + "onnx/*", + "*.h5", + "*.msgpack", + "*.tflite", + "coreml/*", + "*.mlpackage/*", + "*.mlmodel", + "*.gguf", + # Training / checkpoint formats from_pretrained never reads. + "*.pt", + "*.pth", + "*.ckpt", + "optimizer.*", + "scheduler.*", + "rng_state*", + "trainer_state.json", + "events.out.tfevents*", + "checkpoint-*/*", +) + + +# Repo-root tokenizer / config / processor files from_pretrained reads from root even when weights +# load from a subfolder. Exact names (no wildcard) so they match only root-level files. +_ROOT_AUX_PREFETCH_PATTERNS = ( + "config.json", + "generation_config.json", + "tokenizer_config.json", + "tokenizer.json", + "tokenizer.model", + "special_tokens_map.json", + "added_tokens.json", + "vocab.json", + "vocab.txt", + "merges.txt", + "spiece.model", + # More VOCAB_FILES_NAMES the slow tokenizer may fetch (DeBERTa-v2, Whisper, Mistral, XLM-R/mBART, Marian, FSMT/XLM, GPT-2). + "spm.model", + "normalizer.json", + "tokenizer.model.v3", + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + # More VOCAB_FILES_NAMES (RemBERT, FSMT) a distinct-tokenizer-repo warm must cache too. + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + "chat_template.jinja", + "chat_template.json", + # chat_template="" fetches additional_chat_templates/.jinja. + "additional_chat_templates/*.jinja", + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", # Qwen2.5-VL-style video processors + # trust_remote_code auto_map can name any module, so warm every *.py (tiny; none in a non-remote repo). + "*.py", + "*.tiktoken", # tiktoken vocab (e.g. Qwen's qwen.tiktoken) +) + + +# Files a PEFT adapter load reads: config + weights (glob covers sharded adapters). Any merged +# full-model weights the repo also ships match none of these. +_ADAPTER_PREFETCH_PATTERNS = ( + "adapter_config.json", + "adapter_model*", +) + + +# Weight files in a SUBDIRECTORY. A bare root load reads only root weights, so ignoring these drops +# alternate-precision/experimental dirs (fp16/, experimental/). "*/*" spans "/" (HF fnmatch), so nested +# weights match while root "model.safetensors" is kept. Only applied when weights_at_root (diffusion +# keeps weights in subfolders). +_SUBDIR_WEIGHT_IGNORE_PATTERNS = ( + "*/*.safetensors", + "*/*.bin", + "*/*.h5", + "*/*.msgpack", + "*/*.pt", + "*/*.pth", +) + + +def _in_requested_load_scope(filename, subfolder): + """True if *filename* is in the location being loaded (*subfolder*, else root). Scopes the ".bin is + redundant when safetensors exist" test so a .bin-only subfolder keeps its .bin.""" + filename = filename.replace("\\", "/") + if isinstance(subfolder, str) and subfolder.strip("/"): + return filename.startswith(subfolder.strip("/") + "/") + return "/" not in filename # root load: no directory component + + +# .safetensors training-state files that are NOT model weights (e.g. optimizer.safetensors next to a +# real pytorch_model.bin); counting them as "model safetensors present" would drop the needed .bin. +_NON_MODEL_WEIGHT_STEMS = frozenset( + { + "optimizer", + "scheduler", + "scaler", + "rng_state", + "training_args", + } +) + + +def _is_model_weight_safetensors(filename): + """True if *filename* is a model-weights safetensors, not a PEFT adapter/sidecar + (adapter_model.safetensors) or trainer-state (optimizer.safetensors). Only a real one proves the + .bin redundant; counting a sidecar would wrongly drop the needed .bin (fetched then without Xet fallback).""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + if not name.endswith((".safetensors", ".safetensors.index.json")): + return False + if name.startswith("adapter_"): + return False + # Stem before first dot: "optimizer.safetensors" -> "optimizer" (real shards kept); rng_state via prefix. + stem = name.split(".", 1)[0].lower() + if stem in _NON_MODEL_WEIGHT_STEMS or stem.startswith("rng_state"): + return False + return True + + +def _is_canonical_variant_model_weight_safetensors(filename, variant): + """True for a canonical model-weights safetensors carrying the requested *variant*, in the forms + transformers reads (single, either numbered-shard layout, or the index). Strict (base must be + "model"): a sidecar like consolidated..safetensors does not prove the variant .bin redundant.""" + base = filename.replace("\\", "/").rsplit("/", 1)[-1] + v = re.escape(variant) + return bool( + re.match( + rf"^(?:model\.{v}\.safetensors" + rf"|model\.{v}-\d{{5}}-of-\d{{5}}\.safetensors" + rf"|model-\d{{5}}-of-\d{{5}}\.{v}\.safetensors" + rf"|model\.safetensors\.index\.{v}\.json)$", + base, + ) + ) + + +_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE = re.compile( + r"^(?:model\.safetensors|model-\d{5}-of-\d{5}\.safetensors|model\.safetensors\.index\.json)$" +) + + +def _is_canonical_model_weight_safetensors(filename): + """True for a canonical (non-variant) model-weights safetensors a default load reads (model.safetensors, + a numbered shard, or the index). Strict: an unrecognized name keeps both formats, so a variant-only + safetensors + pytorch_model.bin repo never has its .bin dropped for a no-variant load.""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + return bool(_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE.match(name)) + + +def _adapter_repo_has_safetensors( + model_name, + *, + token = None, + revision = None, +): + """Best-effort: does the adapter repo ship a root safetensors adapter weight (making the .bin + redundant)? Scoped to root adapter_model* files; any failure returns False.""" + try: + from huggingface_hub import HfApi + siblings = HfApi().model_info(model_name, revision = revision, token = token).siblings or [] + return any( + "/" not in sibling.rfilename.replace("\\", "/") # root only + and sibling.rfilename.startswith("adapter_model") + and sibling.rfilename.endswith(".safetensors") + for sibling in siblings + ) + except Exception: + return False + + +def _prefetch_ignore_patterns( + model_name, + *, + token = None, + revision = None, + subfolder = None, + use_safetensors = None, + from_tf = False, + from_flax = False, + variant = None, + weights_at_root = False, +): + """ignore_patterns for the prewarm snapshot: the static skip list, minus the checkpoint guard when + loading from a checkpoint-* subfolder, minus the weight format the load will not read. use_safetensors + is a format allowlist (True -> skip *.bin, False -> skip *.safetensors); auto (None) skips *.bin only + when in-scope safetensors are shipped. from_tf/from_flax keep *.h5/*.msgpack. + + Suppressed for a whole multi-component snapshot (weights_at_root=False, no subfolder: ST/diffusers + repos with per-subfolder weights, each in its own format), since "*" spans "/" so dropping "*.bin" + would strip a module's only weight.""" + # Keep checkpoint-*/* under a checkpoint-* subfolder; keep *.h5 / *.msgpack under from_tf/flax. + ignore_patterns = [ + pattern + for pattern in _PREFETCH_IGNORE_PATTERNS + if not ( + ( + pattern == "checkpoint-*/*" + and isinstance(subfolder, str) + and subfolder.startswith("checkpoint-") + ) + or (from_tf and pattern == "*.h5") + or (from_flax and pattern == "*.msgpack") + ) + ] + # Drop the format the load will not read (the other doubles the download); skipped for a whole + # multi-component snapshot (see docstring). + whole_multi_component = not weights_at_root and not ( + isinstance(subfolder, str) and subfolder.strip("/") + ) + if whole_multi_component: + pass + elif from_tf or from_flax: + # TF / Flax loads never read the PyTorch formats; drop safetensors and .bin. + ignore_patterns.extend( + ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + ) + ) + elif use_safetensors is True: + # Explicit safetensors: load never reads .bin (no model_info call needed). + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + elif use_safetensors is False: + # Explicit .bin: load never reads safetensors. + ignore_patterns.extend(("*.safetensors", "*.safetensors.index.json")) + else: + # Auto: skip .bin only once in-scope safetensors are confirmed (best-effort; any failure keeps both). + try: + from huggingface_hub import HfApi + + siblings = ( + HfApi() + .model_info( + model_name, + revision = revision, + token = token, + ) + .siblings + or [] + ) + # Count only in-scope model-weights safetensors (not adapters/sidecars): variant-matching if + # a variant is requested, else canonical, proving the .bin redundant. + has_safetensors = any( + _is_model_weight_safetensors(sibling.rfilename) + and _in_requested_load_scope(sibling.rfilename, subfolder) + and ( + _is_canonical_variant_model_weight_safetensors(sibling.rfilename, variant) + if variant + else _is_canonical_model_weight_safetensors(sibling.rfilename) + ) + for sibling in siblings + ) + if has_safetensors: + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + except Exception: + pass + return ignore_patterns + + +def maybe_prefetch_hf_snapshot( + model_name, + token = None, + *, + revision = None, + cache_dir = None, + local_files_only = False, + fast_inference = False, + subfolder = None, + force_download = False, + use_safetensors = None, + from_tf = False, + from_flax = False, + tokenizer_only = False, + adapter_only = False, + weights_at_root = False, + variant = None, + gguf_file = None, +): + """Warm the HF cache for a remote repo before the in-process load. + + Xet can hang on a blob with no progress or exception, and a blocked native Xet thread cannot be + killed in-process. So pull the snapshot first in a killable subprocess that falls back Xet -> HTTP + on a stall (unsloth_zoo.hf_xet_fallback), making from_pretrained a cache hit. + + Returns True iff warmed (caller can clear force_download), else False (skipped: local/offline/ + local_files_only/fast_inference/old unsloth_zoo, or failed). Only a both-transports-stalled + DownloadStallError is raised; other failures are left for from_pretrained to surface. + """ + try: + from unsloth_zoo.hf_xet_fallback import ( + snapshot_download_with_xet_fallback, + DownloadStallError, + ) + except Exception: + return False # older unsloth_zoo without the helper: load normally + + if not isinstance(model_name, str) or not model_name: + return False + # Local path: nothing to download. Expand ~ first (os.path.exists does not). + model_path = os.path.expanduser(model_name) + if os.path.isdir(model_path) or os.path.exists(model_path): + return False + # Looks local but not yet on disk (e.g. an uncreated output dir): not a Hub repo id, so leave it + # for from_pretrained rather than download it. + if ( + os.path.isabs(model_path) + or model_name.startswith(("~", "./", "../", ".\\", "..\\")) + or "\\" in model_name + ): + return False + if local_files_only: # cache-only: never reach out + return False + if any( + os.environ.get(flag, "0").lower() in ("1", "true", "yes", "on") + for flag in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ): + return False + if fast_inference: # vLLM has its own download path + return False + + # tokenizer-only / adapter-only warms allow-list exact files below, so the weight-format ignore + # list (and its auto-branch model_info call) is skipped. + ignore_patterns = ( + None + if tokenizer_only or adapter_only or gguf_file + else _prefetch_ignore_patterns( + model_name, + token = token, + revision = revision, + subfolder = subfolder, + use_safetensors = use_safetensors, + from_tf = from_tf, + from_flax = from_flax, + variant = variant, + weights_at_root = weights_at_root, + ) + ) + # Narrow the warm to what the load reads (skip extra checkpoints/precisions); every branch still warms + # root tokenizer/config/custom-code so those never fall in-process. + allow_patterns = None + if gguf_file: + # gguf_file=NAME reads exactly that GGUF, but the static ignore list drops *.gguf; so warm just + # that file (plus root aux), under / if set. + _gguf_path = ( + f"{subfolder.strip('/')}/{gguf_file}" + if isinstance(subfolder, str) and subfolder.strip("/") + else gguf_file + ) + allow_patterns = [_gguf_path, *_ROOT_AUX_PREFETCH_PATTERNS] + elif tokenizer_only: + # A distinct tokenizer repo: warm only tokenizer / config / vocab files, never its weights. + allow_patterns = list(_ROOT_AUX_PREFETCH_PATTERNS) + elif adapter_only: + # A PEFT adapter load reads only adapter_config.json + adapter_model.* (plus root aux), not any + # merged weights the repo may also publish. + allow_patterns = [*_ADAPTER_PREFETCH_PATTERNS, *_ROOT_AUX_PREFETCH_PATTERNS] + # PeftModel reads one format (safetensors when present): explicit use_safetensors wins, else + # prefer safetensors when shipped (best-effort; any failure keeps both). + if use_safetensors is False: + ignore_patterns = [ + "adapter_model*.safetensors", + "adapter_model*.safetensors.index.json", + ] + elif use_safetensors is True or _adapter_repo_has_safetensors( + model_name, token = token, revision = revision + ): + ignore_patterns = ["adapter_model*.bin", "adapter_model*.bin.index.json"] + elif isinstance(subfolder, str) and subfolder.strip("/"): + # subfolder=X: load resolves every weight under X/, so warm that subfolder (plus root aux). + allow_patterns = [f"{subfolder.strip('/')}/*", *_ROOT_AUX_PREFETCH_PATTERNS] + elif weights_at_root: + # A bare load reads only root weights: drop subdir weights (fp16/, checkpoint dirs) while keeping + # subdir configs. Diffusion leaves weights_at_root False. + ignore_patterns = [*(ignore_patterns or []), *_SUBDIR_WEIGHT_IGNORE_PATTERNS] + try: + snapshot_download_with_xet_fallback( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + variant = variant, + ) + return True + except DownloadStallError: + # Both transports stalled: surface a clear network error, not a silent in-process hang. + raise + except Exception as exception: + logger.warning_once( + f"Unsloth: Could not pre-download {model_name} " + f"({type(exception).__name__}: {exception}); continuing with the normal load." + ) + return False + + # Ignore logging messages class HideLoggingMessage(logging.Filter): __slots__ = ("text",) @@ -3507,8 +3972,25 @@ def _moe_target_set_from_string(target_modules: str) -> set[str]: return {target_modules} is_regex = re.search(r"[*+?()[\]{}|\\^$]", target_modules) is not None - targets_mlp = "mlp" in target_modules or "ffn" in target_modules - if is_regex and "proj" in target_modules and targets_mlp: + # Key detection on the mlp/ffn/experts path segment (absent from an + # attention-only regex), never on q/k/v/o leaves alone. + targets_mlp_path = any( + tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts") + ) + if not is_regex or not targets_mlp_path: + return set() + # Explicit expert leaves scope the target set to exactly those leaves. + named = {name for name in _MOE_BROAD_MLP_TARGETS if name in target_modules} + if named: + return named + # A generic projection under an mlp path (e.g. ".*mlp.*proj"): any proj + # occurrence that is not an attention leaf name. + if re.search(r"(? Optional[List[str return None +def _select_moe_detection_targets( + original_target_modules, + scoped_target_modules, + finetune_mlp_modules = True, + finetune_language_layers = True, +): + """Pick what get_moe_target_parameters keys expert detection on. + + Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + attention-only request is not pushed into the experts by get_peft_regex's + ``mlp|feed_forward|ffn|dense`` component block (which the string fallback + cannot tell apart from a fused-expert auto regex). + + But only when the MLP and language families are BOTH still in scope. If the + caller scoped MLP or language OFF (``finetune_mlp_modules=False`` or + ``finetune_language_layers=False``) the scoped regex already drops the MoE + experts, and reusing the original list -- which may still name gate/up/down + leaves -- would wrongly re-introduce them. In that case honor the scoped + result so the frozen-MLP / vision-only request is respected. + """ + if original_target_modules is not None and finetune_mlp_modules and finetune_language_layers: + return original_target_modules + return scoped_target_modules + + def make_fast_generate_wrapper(original_generate): """ Creates a wrapper around model.generate that checks for incorrect diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index 0b7f3ab973..cb367d451e 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, ) try: @@ -151,6 +152,9 @@ def CohereAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -161,6 +165,7 @@ def CohereAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py index 12596b432e..955bf55987 100644 --- a/unsloth/models/diffusion.py +++ b/unsloth/models/diffusion.py @@ -24,7 +24,7 @@ import os import torch from transformers import AutoConfig, AutoProcessor, AutoTokenizer -from ._utils import is_bfloat16_supported +from ._utils import is_bfloat16_supported, maybe_prefetch_hf_snapshot from .llama import logger __all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"] @@ -79,7 +79,14 @@ def _resolve_diffusion_model_class(config): ) -def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only): +def _load_diffusion_config( + model_name, + token, + trust_remote_code, + revision, + local_files_only, + cache_dir = None, +): """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4`` classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the type/arch names in-memory, and rebuild.""" @@ -90,6 +97,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except ValueError as e: if "diffusion_gemma" not in str(e): @@ -103,6 +111,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local token = token, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) with open(cfg_path, encoding = "utf-8") as f: cd = json.load(f) @@ -152,12 +161,16 @@ class FastDiffusionModel: os.environ.get("HF_HUB_OFFLINE", "0") == "1" or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1" ) + + cache_dir = kwargs.get("cache_dir") + config = _load_diffusion_config( model_name, token, trust_remote_code, revision, local_files_only, + cache_dir = cache_dir, ) model_type = getattr(config, "model_type", None) if not is_diffusion_model_type(model_type): @@ -168,6 +181,21 @@ class FastDiffusionModel: model_cls = _resolve_diffusion_model_class(config) + # Prefetch the whole repo root so the weight load is a cache hit. No subfolder: the pipeline + # loads every component subfolder, so narrowing would leave unet/vae/text_encoder to Xet. + maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + local_files_only = local_files_only, + fast_inference = False, + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + # Forward variant (e.g. "fp16") so the warm keeps variant weights. + variant = kwargs.get("variant"), + ) + load_kwargs = dict( dtype = dtype, device_map = device_map, @@ -176,7 +204,14 @@ class FastDiffusionModel: attn_implementation = attn_implementation, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) + # Match the load's weight format to the warm (None/auto already matches). + if kwargs.get("use_safetensors") is not None: + load_kwargs["use_safetensors"] = kwargs["use_safetensors"] + # Forward variant to the real load so it reads the warmed variant weights. + if kwargs.get("variant") is not None: + load_kwargs["variant"] = kwargs["variant"] # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision. @@ -222,6 +257,7 @@ class FastDiffusionModel: trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except Exception: tokenizer = AutoTokenizer.from_pretrained( @@ -230,6 +266,7 @@ class FastDiffusionModel: trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) return model, tokenizer diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 68b9ebe22f..4a0531db78 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .gemma import ( @@ -168,6 +169,11 @@ def Gemma2Attention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. gemma2 is + # sliding-window and softcapped: the engage gate caps spans at the window and + # excludes softcap models entirely, so PG never engages here. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -179,6 +185,7 @@ def Gemma2Attention_fast_forward( attention_mask = attention_mask, causal_mask = causal_mask, sliding_window = sliding_window, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index f5b0f57aa6..4dedf642eb 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .llama import ( @@ -159,6 +160,9 @@ def GraniteAttention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -169,6 +173,7 @@ def GraniteAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index bb7289dfa8..c25a031b82 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -39,6 +39,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from torch.nn.functional import scaled_dot_product_attention from transformers import __version__ as transformers_version @@ -738,6 +739,10 @@ def LlamaAttention_fast_forward( flash_dense_kwargs = {"causal": True}, flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True}, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward (same route + # as packed_seq_lengths); misuse (KV cache / padding mask) raises. None => byte-identical + # default. Reuse of this forward also carries the branch to qwen2 & gemma. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -748,6 +753,7 @@ def LlamaAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = config, context = context, Q = Q, K = K, V = V) @@ -895,8 +901,10 @@ def LlamaModel_fast_forward( seq_length_with_past = seq_length # Fix out of bounds tokenization unless we were given packed metadata - allow_overlength = getattr(self, "_unsloth_allow_packed_overlength", False) or ( - "packed_seq_lengths" in kwargs + allow_overlength = ( + getattr(self, "_unsloth_allow_packed_overlength", False) + or ("packed_seq_lengths" in kwargs) + or ("prefix_seg_info" in kwargs and kwargs["prefix_seg_info"] is not None) ) if hasattr(self, "max_seq_length") and not allow_overlength: if seq_length > self.max_seq_length: @@ -1748,7 +1756,6 @@ class LlamaRotaryEmbedding(torch.nn.Module): # Base-class-from-config path (modern transformers): derive inv_freq like # transformers so config.rope_scaling is not dropped (#2405). Scaled # subclasses are excluded to avoid double-scaling. - config_inv_freq = None if config is not None: # [TODO] Hack to pass in config - need to remove later base = _get_rope_theta(config, default = base) @@ -1761,32 +1768,17 @@ class LlamaRotaryEmbedding(torch.nn.Module): device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is not None and type(self) is LlamaRotaryEmbedding: - config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq( - config, - rope_scaling, - ) - self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base + # Kept so the v5 rope repair can rebuild the scaled inv_freq (#2405). + self._unsloth_rope_config = config # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this self.current_rope_size = min(4 * 8192, self.max_position_embeddings) self.multi_gpu_cos_cached = [None] * DEVICE_COUNT self.multi_gpu_sin_cached = [None] * DEVICE_COUNT - if config_inv_freq is not None: - inv_freq = config_inv_freq # already scaled; skip subclass scaling - else: - # Normal Llama-3 RoPE - inv_freq = 1.0 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim - ) - ) - inv_freq = self._apply_inv_freq_scaling(inv_freq) + inv_freq = self._unsloth_recompute_inv_freq() self.register_buffer("inv_freq", inv_freq, persistent = False) # Build here to make `torch.jit.trace` work. @@ -1809,6 +1801,25 @@ class LlamaRotaryEmbedding(torch.nn.Module): """Override to apply custom inv_freq scaling (e.g., extended RoPE).""" return inv_freq + def _unsloth_recompute_inv_freq(self): + # Config scaling (llama3/yarn) first, else vanilla + subclass scaling. + # Shared by __init__ and the v5 rope repair so they cannot diverge. + config = getattr(self, "_unsloth_rope_config", None) + config_inv_freq = None + rope_scaling = getattr(config, "rope_scaling", None) if config is not None else None + if rope_scaling is not None and type(self) is LlamaRotaryEmbedding: + config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq( + config, + rope_scaling, + ) + if config_inv_freq is not None: + return config_inv_freq + inv_freq = 1.0 / ( + self.base + ** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim) + ) + return self._apply_inv_freq_scaling(inv_freq) + def _apply_time_scaling(self, t): """Override to apply custom time scaling (e.g., linear scaling).""" return t @@ -2420,6 +2431,73 @@ class FastLlamaModel: preferred_attn_impl = resolve_attention_implementation(model_function, model_config) + # Prefetch the repo (killable child) so the weight load is a cache hit. Runs after the + # AutoConfig/model-class check so an unsupported repo fails on its small config fetch. No + # revision: the load resolves model_name (maybe a remapped prequant repo) on its default branch. + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + # Skip the warm only for a real vLLM load; a num_labels classification load still goes + # in-process below, so it must be warmed even under fast_inference. + fast_inference = fast_inference and num_labels is None, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Tokenizer always loads in-process. Resolve the cache_dir the tokenizer load will actually + # use, mirroring load_correct_tokenizer: without an explicit cache_dir, Colab/Kaggle route to + # a special tokenizer cache (huggingface_tokenizers_cache / Kaggle tmp), NOT the HF-default + # cache the base snapshot warmed. So the base warm does not cover the tokenizer there. + from ..tokenizer_utils import ( + IS_COLAB_ENVIRONMENT, + IS_KAGGLE_ENVIRONMENT, + KAGGLE_TMP, + ) + + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _tokenizer_cache_dir = kwargs.get("cache_dir") + if _tokenizer_cache_dir is None: + if IS_COLAB_ENVIRONMENT: + _tokenizer_cache_dir = "huggingface_tokenizers_cache" + elif IS_KAGGLE_ENVIRONMENT: + _tokenizer_cache_dir = os.path.join(KAGGLE_TMP, "huggingface_tokenizers_cache") + # Warm the tokenizer repo into the cache the load will use whenever the base warm did not + # cover it: a distinct tokenizer repo, fast_inference (base warm skipped), or a tokenizer + # cache_dir that differs from the base-warm cache_dir (Colab/Kaggle special cache). + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and ( + _tokenizer_repo != model_name + or fast_inference + or _tokenizer_cache_dir != kwargs.get("cache_dir") + ) + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = _tokenizer_cache_dir, + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + has_rope_scaling = False try: with open(inspect.getfile(model_function), "r", encoding = "utf-8") as file: @@ -2672,6 +2750,10 @@ class FastLlamaModel: # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # Route the tokenizer load to the custom cache_dir the prefetch warmed. + _tokenizer_cache_kwargs = {} + if kwargs.get("cache_dir") is not None: + _tokenizer_cache_kwargs["cache_dir"] = kwargs["cache_dir"] tokenizer = load_correct_tokenizer( tokenizer_name = tokenizer_name, model_max_length = max_position_embeddings, @@ -2679,6 +2761,7 @@ class FastLlamaModel: token = token, trust_remote_code = trust_remote_code, fix_tokenizer = fix_tokenizer, + **_tokenizer_cache_kwargs, ) model, tokenizer = patch_tokenizer(model, tokenizer) @@ -2805,6 +2888,7 @@ class FastLlamaModel: model_max_length = max_position_embeddings, padding_side = "right", token = token, + cache_dir = kwargs.get("cache_dir"), ) patch_saving_functions(tokenizer) @@ -3743,4 +3827,17 @@ class FastLlamaModel: from .rl import PatchFastRL +# Auto-enable grouped-GEMM MoE (tf<5 ModuleList experts) on built / PEFT'd models. Wrap the +# loader leaves before PatchFastRL so downstream patchers see the wrapped versions. Guarded. +try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe + FastLlamaModel.from_pretrained = staticmethod( + wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained) + ) + FastLlamaModel.get_peft_model = staticmethod( + wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model) + ) +except Exception: + pass + PatchFastRL(FastLanguageModel = FastLlamaModel) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 562afdd645..13342157b0 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -21,6 +21,10 @@ from ._utils import ( USE_MODELSCOPE, get_transformers_model_type, hf_login, + # Single source of truth is _utils.py; re-exported here so callers doing + # `from unsloth.models.loader import DISABLE_SDPA_MODEL_NAMES` keep working and so + # _is_sdpa_excluded (in _utils) can honor it without a loader -> _utils cycle. + DISABLE_SDPA_MODEL_NAMES, ) from .granite import FastGraniteModel from .llama import FastLlamaModel, logger @@ -106,24 +110,31 @@ from ._utils import ( _is_family_text_decoder, _apply_text_only_key_mapping, set_task_config_attr, + maybe_prefetch_hf_snapshot, ) -# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers -# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working. -# Fallback list mirrors zoo for users who upgrade unsloth without upgrading -# unsloth_zoo (so this module never fails at import). +# Source of truth is unsloth_zoo.model_lists. Re-exported so callers doing +# `from unsloth.models.loader import FORCE_FLOAT32` keep working. The fallback +# list is also unioned in so a newer unsloth still forces float32 for these +# archs when paired with an older unsloth_zoo that predates them (upgrade skew). +_FORCE_FLOAT32_FALLBACK = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) + "gemma3n", + "gemma4", # Gemma4 (gemma4 / gemma4_text): float16 NaNs grad norms in the backward + "glm4_moe", # GLM-4.x MoE (glm4_moe / glm4_moe_lite): float16 NaNs grad norms + "gpt_oss", + "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training + "qwen3_moe", # Qwen3-MoE (Qwen3-30B-A3B): float16 NaNs grad norms in the backward +] try: - from unsloth_zoo import FORCE_FLOAT32 # noqa: F401 + from unsloth_zoo import FORCE_FLOAT32 as _ZOO_FORCE_FLOAT32 + FORCE_FLOAT32 = list(_ZOO_FORCE_FLOAT32) except ImportError: - global FORCE_FLOAT32 - # Forces float32 precision since float16 goes to infinity - FORCE_FLOAT32 = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) - "gemma3n", - "gpt_oss", - "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training - ] + FORCE_FLOAT32 = [] +for _mt in _FORCE_FLOAT32_FALLBACK: + if not any(_mt in _entry for _entry in FORCE_FLOAT32): + FORCE_FLOAT32.append(_mt) global DISABLE_COMPILE_MODEL_NAMES # Must be alphabetically sorted for each entry @@ -195,13 +206,44 @@ DISABLE_COMPILE_MODEL_NAMES = [ "granite,llava_next", # Granite-vision 3 ] -global DISABLE_SDPA_MODEL_NAMES -# Disables some SDPA modules since it's wrong -DISABLE_SDPA_MODEL_NAMES = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore - "gpt_oss", -] +# Architectures with gated-deltanet (linear attention) layers. Unsloth bundles the +# flash-linear-attention Triton kernels (unsloth_zoo/_vendored/fla), so no install is +# needed; transformers uses the much slower pure PyTorch path only when they can't be enabled. +FLA_MODEL_TYPE_PREFIXES = ("qwen3_next", "qwen3_5", "kimi_linear", "olmo_hybrid") +_fla_advised = False + + +def _maybe_advise_fla_install(model_types): + """One-time note when a gated-deltanet model loads without the fast kernels. + + The kernels ship with Unsloth (no install needed); this fires only when they + could not be enabled on this platform (e.g. no CUDA, torch < 2.7 or + triton < 3.3), i.e. exactly when transformers uses the slow pure PyTorch path. + """ + global _fla_advised + if _fla_advised: + return + if model_types is None: + return + if isinstance(model_types, str): + model_types = [model_types] # a lone string would otherwise iterate chars + try: + if not any( + isinstance(t, str) and t.startswith(FLA_MODEL_TYPE_PREFIXES) for t in model_types + ): + return + from transformers.utils.import_utils import is_flash_linear_attention_available + if is_flash_linear_attention_available(): + return # bundled (or user-installed) fast kernels are active + except Exception: + return + _fla_advised = True + print( + "Unsloth: This model uses gated-deltanet linear attention layers. Unsloth\n" + "bundles the flash-linear-attention kernels, but they could not be enabled\n" + "on this setup (they need CUDA with torch >= 2.7 and triton >= 3.3), so\n" + "transformers will use a slower pure PyTorch path." + ) def _fix_rope_inv_freq(model): @@ -227,14 +269,18 @@ def _fix_rope_inv_freq(model): and hasattr(module, "_apply_inv_freq_scaling") and hasattr(module, "multi_gpu_cos_cached") ): - inv_freq = 1.0 / ( - module.base - ** ( - torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float() - / module.dim + if hasattr(module, "_unsloth_recompute_inv_freq"): + # Restore config scaling (llama3/yarn); unscaled here broke v5. + inv_freq = module._unsloth_recompute_inv_freq() + else: + inv_freq = 1.0 / ( + module.base + ** ( + torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float() + / module.dim + ) ) - ) - inv_freq = module._apply_inv_freq_scaling(inv_freq) + inv_freq = module._apply_inv_freq_scaling(inv_freq) module.inv_freq = inv_freq for device_idx in range(len(module.multi_gpu_cos_cached)): if module.multi_gpu_cos_cached[device_idx] is not None: @@ -469,8 +515,10 @@ class FastLanguageModel(FastLlamaModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -628,8 +676,10 @@ class FastLanguageModel(FastLlamaModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -865,6 +915,28 @@ class FastLanguageModel(FastLlamaModel): if is_peft: # From https://github.com/huggingface/peft/issues/184 # Now add PEFT adapters + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] model = PeftModel.from_pretrained( model, old_model_name, @@ -873,9 +945,19 @@ class FastLanguageModel(FastLlamaModel): local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) # Patch it as well! model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing) + # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back + # to the original loop, an attention-only adapter keeps the grouped path. Guarded. + try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import ( + auto_enable_grouped_moe, + ) + auto_enable_grouped_moe(model) + except Exception: + pass # optional speedup; never block model loading # Patch Tiled MLP # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"" @@ -1116,8 +1198,10 @@ class FastModel(FastBaseModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -1263,6 +1347,7 @@ class FastModel(FastBaseModel): trust_remote_code = trust_remote_code, ) model_types_all = ",".join(model_types) + "," + _maybe_advise_fla_install(model_types) # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ---- # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's @@ -1474,8 +1559,10 @@ class FastModel(FastBaseModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -1790,6 +1877,28 @@ class FastModel(FastBaseModel): _LoraModel._create_and_replace = _patched_car + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] try: model = PeftModel.from_pretrained( model, @@ -1799,6 +1908,7 @@ class FastModel(FastBaseModel): local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) finally: # Always restore original PEFT method, even if loading fails @@ -1809,6 +1919,15 @@ class FastModel(FastBaseModel): model = FastBaseModel.post_patch_model( model, use_gradient_checkpointing, trust_remote_code = trust_remote_code ) + # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back + # to the original loop, an attention-only adapter keeps the grouped path. Guarded. + try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import ( + auto_enable_grouped_moe, + ) + auto_enable_grouped_moe(model) + except Exception: + pass # optional speedup; never block model loading # Apply QAT if specified if qat_scheme is not None: diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index df2a4de5bd..4350565fe2 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -27,6 +27,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -124,6 +125,9 @@ def MistralAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -134,6 +138,7 @@ def MistralAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) @@ -161,7 +166,13 @@ def MistralForCausalLM_fast_forward( *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - if causal_mask is None and past_key_values is None: + # PrefixGrouper brings its own mask: a synthesized causal attention_mask would trip + # resolve_prefix_seg_info on the no-xFormers path and force a fallback. + if ( + causal_mask is None + and past_key_values is None + and kwargs.get("prefix_seg_info", None) is None + ): bsz, q_len = input_ids.shape sliding_window = getattr(self.config, "sliding_window", None) diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index e28e72d3ea..0d05a2d538 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -146,6 +147,9 @@ def Qwen3Attention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -156,6 +160,7 @@ def Qwen3Attention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 3be614cf4a..098950de08 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -29,7 +29,9 @@ from collections import defaultdict from unsloth_zoo.rl_replacements import ( RL_REPLACEMENTS, left_pack_padding, + create_completion_attention_mask, chunked_selective_log_softmax, + chunked_hidden_states_selective_log_softmax, _unsloth_get_mm_token_id, _unsloth_fix_mm_token_type_ids, ) @@ -48,7 +50,41 @@ from ..device_type import ( ALLOW_PREQUANTIZED_MODELS, ) import textwrap -from ._utils import _get_inference_mode_context_manager +from ._utils import _get_inference_mode_context_manager, UNSLOTH_ENABLE_LOGGING + +# One-time GRPO sequence-packing gates; mirrored into the generated trainer cache via RL_PRE_ITEMS. +UNSLOTH_GRPO_SEQ_PACKING_ON = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "1").lower() not in ( + "0", + "false", + "no", + "off", +) +# Packing needs zoo#840's masked-column guard in grpo_compute_loss (installed zoo is fixed per-process). +try: + UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = "torch.where(_keep, new" in inspect.getsource( + RL_REPLACEMENTS["grpo_compute_loss"] + ) +except Exception: + UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False +# One-time PrefixGrouper gate; any import failure degrades to "PrefixGrouper off". +_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None +UNSLOTH_GRPO_PREFIX_GROUPER_ON = os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ( + "0", + "false", + "no", + "off", +) +if UNSLOTH_GRPO_PREFIX_GROUPER_ON: + try: + from ..utils.prefix_grouper import ( + build_group_layout as _pg_build_layout, + prefix_grouper_enabled as _pg_enabled_fn, + verify_on as _pg_verify_on, + tol_ok as _pg_tol_ok, + TOL_KILL as _PG_TOL_KILL, + ) + except Exception: + UNSLOTH_GRPO_PREFIX_GROUPER_ON = False RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) @@ -1359,6 +1395,439 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): ) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- + # One varlen [1, sum L] forward replaces the padded [B, Lmax] loop (also fixes the + # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T + # grows; falls back if a backend ignores packed_seq_lengths. + logprobs = None + + # ---- PrefixGrouper (GRPO shared-prompt dedup; default ON, exact + self-verified) ---- + # G completions per prompt share the prefix; the packed path forwards it G times, + # PrefixGrouper stores it once (FlexAttention shared-prefix mask), cutting the trunk + # forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER (needs + # seq-packing), tok_r auto-gate, and first-use self-verify vs the packed path + # (mismatch => fall back + mark unsafe), so a mask/isolation regression cannot ship + # silently. When off / ungrouped / unverified, the packed path below runs as before. + _pg_result = None + _pg_use = False + _pg_skip_pk = False # once a shape is PG-verified, skip the full-row forward + _pg_forward_fn = None # deferred PG forward (runs at the verify site below) + _pg_num_gen = getattr(self, "num_generations", None) + # Env gate hoisted to module level (mirrored via RL_PRE_ITEMS). Skip PG under vLLM + # (fast_inference=True): the rollout dominates the step, so PG saves little and its + # first-use self-verify is net overhead. + _pg_engage = ( + UNSLOTH_GRPO_PREFIX_GROUPER_ON + and not getattr(self, "use_vllm", False) + and not getattr(unwrapped_model, "_unsloth_prefix_grouper_nograd_disabled", False) + ) + if _pg_engage: + try: + # Skip softcap models (the flex kernel never applies attn_logit_softcapping) + # and hybrid SSM / MoE models: only the threaded attention forwards get the + # shared-prefix isolation, so a Mamba or MoE decoder that does not forward + # prefix_seg_info would leak suffixes across completions. PG also rides on + # sequence packing, so it needs the same zoo masked-column guard. + _pg_cfg = getattr(unwrapped_model, "config", None) + _pg_engage = ( + _pg_enabled_fn() + and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD + and pixel_values is None + and token_type_ids is None + and mm_token_type_ids is None + and _pg_num_gen is not None + and _pg_num_gen >= 2 + and not getattr(_pg_cfg, "attn_logit_softcapping", None) + # normal backends apply config.attention_dropout in training; the flex + # path is deterministic, so skip PG when it is set. + and not getattr(_pg_cfg, "attention_dropout", 0) + and not any( + getattr(_pg_cfg, _pg_a, None) is not None + for _pg_a in ( + "mamba_d_ssm", + "mamba_d_state", + "mamba_expand", + "num_experts", + "num_local_experts", + "n_routed_experts", + "moe_intermediate_size", + ) + ) + ) + except Exception: + _pg_engage = False + if _pg_engage: + try: + _pg_pad = self.processing_class.pad_token_id + # cap the PG span (P+max(R)) at the sliding window, like the packed _pk_sw guard. + _pg_sw = getattr( + getattr(unwrapped_model, "config", None), "sliding_window", None + ) + if not (isinstance(_pg_sw, int) and _pg_sw > 0): + _pg_sw = None + _pg_layout = _pg_build_layout( + input_ids, + logits_to_keep, + _pg_pad, + _pg_num_gen, + left_pad_tokens_per_prompt, + max_segment_cap = _pg_sw, + ) + _pg_unsafe = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_unsafe is None: + _pg_unsafe = set() + if _pg_layout is not None and _pg_layout.signature not in _pg_unsafe: + _pg_sig = _pg_layout.signature + _pg_verified = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if _pg_verified is None: + _pg_verified = set() + _pg_chunks = max(1, total_rows * multiplier) + + def _pg_run_forward(_pg_layout = _pg_layout, _pg_chunks = _pg_chunks): + with _get_inference_mode_context_manager(model): + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + _pg_hidden = unwrapped_model( + input_ids = _pg_layout.flat_ids, + position_ids = _pg_layout.position_ids, + prefix_seg_info = _pg_layout.prefix_seg_info, + use_cache = False, + ).logits + _pg_r = _pg_layout.extract_logps( + _pg_hidden, + lm_head, + chunked_hidden_states_selective_log_softmax, + _pg_chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) + _pg_hidden = None # release before any verify forward + device_synchronize() + # clip to the loss window [B, logits_to_keep+max_left_pad] + _pg_w = logits_to_keep + max_left_pad + if _pg_r.shape[1] > _pg_w: + _pg_r = _pg_r[:, -_pg_w:] + return _pg_r + + # trust only within the verified envelope: re-verify when T or the + # longest segment grows, like the packed path + _pg_T = int(_pg_layout.flat_ids.shape[1]) + _pg_maxseg = int(_pg_layout.position_ids.max()) + 1 + _pg_env = ( + _pg_verified.get(_pg_sig) if isinstance(_pg_verified, dict) else None + ) + if (not _pg_verify_on()) or ( + _pg_env is not None and _pg_T <= _pg_env[0] and _pg_maxseg <= _pg_env[1] + ): + # trusted shape: run PG now and skip the full-row forward below + _pg_result = _pg_run_forward() + _pg_use = True + _pg_skip_pk = True + else: + # unverified shape: defer the forward until the packed reference + # exists (verify site below), so a declined packed path never wastes + # a whole-batch PG forward + _pg_forward_fn = _pg_run_forward + except Exception as _pg_err: + _pg_result = None + _pg_use = False + _pg_skip_pk = False + _pg_forward_fn = None + # A FlexAttention/Triton compile failure or OOM here is GPU-wide, not + # layout-specific, so retrying the same PG forward every step just re-pays + # the failure. Persistently disable PG (mirrors the seq-packing handler + # setting _unsloth_seq_packing_nograd_ok = False); the packed/padded path + # below still produces the exact result. + unwrapped_model._unsloth_prefix_grouper_nograd_disabled = True + if isinstance(_pg_err, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) disabled (fell back to packed): {_pg_err!r}", + flush = True, + ) + + # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- + # One varlen [1, sum L] block-diagonal forward replaces the padded [B, Lmax] loop + # (exact per-row result; also fixes the padded path's left-pad RoPE error). + # Self-verified vs the per-row forward, re-checked as T grows; falls back if a + # backend ignores packed_seq_lengths. lm_head runs on completion positions only. + _pk_result = None + _pk_use = False + _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON + # Without zoo#840's masked-column guard, zeroed prompt/pad columns turn NaN in exp(). + _pk_enabled = _pk_enabled and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD + _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None) + if ( + _pk_enabled + and not _pg_skip_pk + and pixel_values is None + and token_type_ids is None + and mm_token_type_ids is None + and _pk_ok is not False + ): + try: + _pk_pad = self.processing_class.pad_token_id + _pk_keep = input_ids != _pk_pad + _pk_len = _pk_keep.sum(dim = 1) + _pk_len_cpu = _pk_len.tolist() # single GPU->CPU sync, reused below + _pk_nz_cpu = [_n for _n in _pk_len_cpu if _n > 0] + _pk_flat = input_ids[_pk_keep].unsqueeze(0) + _pk_T = _pk_flat.shape[1] + _pk_L = input_ids.shape[1] + _pk_W = logits_to_keep + max_left_pad + _pk_maxseg = max(_pk_nz_cpu) if _pk_nz_cpu else 0 + # sliding-window models lose the per-sequence local window in a packed stream + _pk_sw = getattr( + getattr(unwrapped_model, "config", None), "sliding_window", None + ) + _pk_sw_ok = not (isinstance(_pk_sw, int) and _pk_sw > 0 and _pk_maxseg > _pk_sw) + # per-row completion mask (same as the loss); prompt-only rows count as inactive + _pk_cmask = create_completion_attention_mask( + input_ids[:, -_pk_W:], left_pad_tokens_per_prompt, max_left_pad, _pk_pad + ) + _pk_active = int(_pk_cmask.any(dim = 1).sum()) + # skip the packed forward entirely at known-unsafe lengths (avoids a wasted pass / OOM) + _pk_unsafe = getattr( + unwrapped_model, "_unsloth_seq_packing_nograd_unsafe_T", None + ) + # cap the flattened forward at one padded [batch_size, seq_len] mini-batch's + # token budget; anything larger uses the chunked padded loop + _pk_cap = batch_size * seq_len + if ( + _pk_T >= 2 + and _pk_T <= _pk_cap + and len(_pk_nz_cpu) > 0 + and _pk_sw_ok + and not (_pk_unsafe is not None and _pk_T >= _pk_unsafe) + and (_pk_ok is True or _pk_active >= 2) + ): + # reset 0-based position_ids per segment + _pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0) + _pk_chunks = max(1, total_rows * multiplier) + _pk_nz_idx = _pk_keep.nonzero( + as_tuple = False + ) # [T, 2] = (row, col), row-major + _pk_within = _pk_nz_idx[1:, 0] == _pk_nz_idx[:-1, 0] # [T-1] + # per-row completion start after left-packing (matches create_completion_attention_mask) + _pk_cstart = (_pk_L - logits_to_keep) - left_pad_tokens_per_prompt # [rows] + _pk_ctgt = (_pk_nz_idx[1:, 1] >= _pk_cstart[_pk_nz_idx[1:, 0]]) & _pk_within + with _get_inference_mode_context_manager(model): + with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): + # use_cache=False: a KV cache silently disables varlen packing + _pk_hidden = unwrapped_model( + input_ids = _pk_flat, + position_ids = _pk_pos, + packed_seq_lengths = torch.tensor( + _pk_nz_cpu, dtype = torch.int32, device = input_ids.device + ), + use_cache = False, + ).logits + _pk_sel = chunked_hidden_states_selective_log_softmax( + _pk_hidden[0, :-1, :][_pk_ctgt].unsqueeze(0), + lm_head, + _pk_flat[0, 1:][_pk_ctgt].unsqueeze(0), + _pk_chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] + # GPT-OSS offload race guard (matches the padded loop) + device_synchronize() + # scatter each logprob back to its (row, col) so [:, -_pk_W:] matches padded + _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt] + _pk_result = ( + torch.zeros( + total_rows * _pk_L, + dtype = torch.float32, + device = input_ids.device, + ) + .index_put((_pk_tgt,), _pk_sel.to(torch.float32)) + .view(total_rows, _pk_L)[:, -_pk_W:] + ) + # re-verify when T or the longest segment grows past what was verified + # (a LongRoPE cache switch can change the result) + _pk_vT = int( + getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_T", 0) + ) + _pk_vS = int( + getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_seg", 0) + ) + # debug: hand-edit this condition to force re-verify every step + if _pk_ok is True and _pk_T <= _pk_vT and _pk_maxseg <= _pk_vS: + _pk_use = True # already verified for this shape + else: + # verify against the per-row forward (ground truth) + _pk_ref = torch.zeros_like(_pk_result) + with _get_inference_mode_context_manager(model): + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + for _pk_i in range(total_rows): + _pk_ni = _pk_len_cpu[_pk_i] + if _pk_ni < 2: + continue + _pk_rmask = _pk_keep[_pk_i] + _pk_real = input_ids[_pk_i][_pk_rmask].unsqueeze(0) + _pk_rpos = torch.arange( + _pk_ni, device = input_ids.device + ).unsqueeze(0) + _pk_rh = unwrapped_model( + input_ids = _pk_real, + position_ids = _pk_rpos, + use_cache = False, + ).logits + _pk_rsel = chunked_hidden_states_selective_log_softmax( + _pk_rh[:, :-1, :], + lm_head, + _pk_real[:, 1:], + 1, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] + _pk_rcols = _pk_rmask.nonzero(as_tuple = False).squeeze(1)[ + 1: + ] - (_pk_L - _pk_W) + _pk_rkeep = _pk_rcols >= 0 + _pk_ref[_pk_i, _pk_rcols[_pk_rkeep]] = _pk_rsel[ + _pk_rkeep + ].to(torch.float32) + device_synchronize() + # compare over the loss-mask region only + _pk_cm = _pk_cmask.float() + _pk_diff = float(((_pk_result - _pk_ref).abs() * _pk_cm).max()) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO seq-packing (no-grad) verify: T={_pk_T} maxseg={_pk_maxseg} packed-vs-perrow max|d|={_pk_diff:.4f}", + flush = True, + ) + # kernel-noise floor ~0.25; cross-sample contamination is >= 2.4 + if _pk_diff < 7e-1: + unwrapped_model._unsloth_seq_packing_nograd_ok = True + # widen the trusted shape only when >= 2 completion rows exercised + # cross-sample packing; single-row passes prove nothing + if _pk_active >= 2: + unwrapped_model._unsloth_seq_packing_nograd_verified_T = max( + _pk_vT, _pk_T + ) + unwrapped_model._unsloth_seq_packing_nograd_verified_seg = max( + _pk_vS, _pk_maxseg + ) + _pk_ok = True + _pk_use = True + else: + _pk_use = False + if _pk_diff >= 1.5: + # contamination (attention ignores the packed mask): disable packing + unwrapped_model._unsloth_seq_packing_nograd_ok = False + else: + # likely a length boundary (LongRoPE): mark unsafe, keep smaller shapes + unwrapped_model._unsloth_seq_packing_nograd_unsafe_T = ( + _pk_T if _pk_unsafe is None else min(_pk_unsafe, _pk_T) + ) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO seq-packing (no-grad) fell back at T={_pk_T} (diff={_pk_diff:.3f})", + flush = True, + ) + except Exception as _pk_err: + # any failure: drop intermediates, use the padded loop, do not retry + _pk_hidden = None + _pk_sel = None + _pk_result = None + _pk_use = False + if isinstance(_pk_err, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + unwrapped_model._unsloth_seq_packing_nograd_ok = False + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}", + flush = True, + ) + # ---- PrefixGrouper first-use self-verify (no-grad) ---- + # Compare the untrusted PG result to the full-row packed result (itself verified vs + # per-row) over the completion mask: < tol_ok -> trust the structure; >= TOL_KILL -> + # unsafe forever; borderline -> fall back this shape. + if _pg_forward_fn is not None and not _pg_use: + if _pk_use and _pk_result is not None: + try: + # deferred PG forward, run only now that the packed reference exists + _pg_result = _pg_forward_fn() + _pg_W2 = logits_to_keep + max_left_pad + _pg_cm = create_completion_attention_mask( + input_ids[:, -_pg_W2:], + left_pad_tokens_per_prompt, + max_left_pad, + self.processing_class.pad_token_id, + ).float() + _pg_a = _pg_result[:, -_pg_W2:].float() + _pg_b = _pk_result[:, -_pg_W2:].float() + _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max()) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify: sig={_pg_layout.signature} " + f"shared-prefix vs full-row-packed max|d|={_pg_diff:.4f}", + flush = True, + ) + if _pg_diff < _pg_tol_ok(): + _pg_v = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if not isinstance(_pg_v, dict): + _pg_v = {} + _pg_vT = int(_pg_layout.flat_ids.shape[1]) + _pg_vS = int(_pg_layout.position_ids.max()) + 1 + _pg_old = _pg_v.get(_pg_layout.signature, (0, 0)) + _pg_v[_pg_layout.signature] = ( + max(_pg_vT, _pg_old[0]), + max(_pg_vS, _pg_old[1]), + ) + unwrapped_model._unsloth_prefix_grouper_nograd_verified = _pg_v + _pg_use = True + else: + _pg_u = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_u is None: + _pg_u = set() + if _pg_diff >= _PG_TOL_KILL: + _pg_u.add(_pg_layout.signature) + unwrapped_model._unsloth_prefix_grouper_nograd_unsafe = _pg_u + _pg_use = False + except Exception as _pg_err3: + _pg_result = None + _pg_use = False + if isinstance(_pg_err3, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify failed (fell back to packed): {_pg_err3!r}", + flush = True, + ) + # else: no packed reference (packing off/failed) -> cannot verify; fall back. + + if _pg_use and _pg_result is not None: + logprobs = _pg_result # PrefixGrouper verified/trusted -> skip the loop + zipped_inputs = [] + elif _pk_use and _pk_result is not None: + logprobs = _pk_result # verified -> skip the loop + zipped_inputs = [] + else: + # free packed intermediates before running the padded loop + _pk_hidden = _pk_sel = _pk_result = _pk_ref = None + with _get_inference_mode_context_manager(model): for ( input_ids_chunk, @@ -1443,7 +1912,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # However, it seems that this line does not slow down or disrupt models. device_synchronize() all_logprobs_list.append(logprobs_chunk) - logprobs = torch.cat(all_logprobs_list, dim = 0) + if logprobs is None: # padded fallback when packing was not used + logprobs = torch.cat(all_logprobs_list, dim = 0) entropies = None os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" @@ -1523,6 +1993,34 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss)) RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_get_inference_mode_context_manager)) +# inspect.getsource inlines function bodies but not module imports, so constants the inlined +# grpo functions reference (e.g. UNSLOTH_ENABLE_LOGGING) must be redefined in the generated cache. +RL_PRE_ITEMS["grpo_trainer"].append( + "import os as _unsloth_os\n" + "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n" +) +# Sequence-packing gates, same values as the module-top constants. +RL_PRE_ITEMS["grpo_trainer"].append( + "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n" +) +RL_PRE_ITEMS["grpo_trainer"].append( + "try:\n" + " import inspect as _unsloth_inspect\n" + " from unsloth_zoo.rl_replacements import RL_REPLACEMENTS as _unsloth_zoo_RL\n" + " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = 'torch.where(_keep, new' in _unsloth_inspect.getsource(_unsloth_zoo_RL['grpo_compute_loss'])\n" + "except Exception:\n" + " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n" +) +# PrefixGrouper gate, same shape as the module-top constants. +RL_PRE_ITEMS["grpo_trainer"].append( + "_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None\n" + "UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off')\n" + "if UNSLOTH_GRPO_PREFIX_GROUPER_ON:\n" + " try:\n" + " from unsloth.utils.prefix_grouper import build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL\n" + " except Exception:\n" + " UNSLOTH_GRPO_PREFIX_GROUPER_ON = False\n" +) # Edit _get_per_token_logps to handle mixed precision diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 7e43442bfd..c1172faa94 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -19,6 +19,7 @@ from ._utils import ( SUPPORTS_BFLOAT16, resolve_model_class, resolve_encoder_attention_implementation, + maybe_prefetch_hf_snapshot, ) import inspect import json @@ -541,7 +542,12 @@ class FastSentenceTransformer(FastModel): return transformer_module @staticmethod - def _read_pooling_mode(model_name, token): + def _read_pooling_mode( + model_name, + token, + cache_dir = None, + revision = None, + ): """Read the pooling mode from modules.json, else return "mean".""" try: if os.path.exists(model_name) and os.path.exists( @@ -549,7 +555,13 @@ class FastSentenceTransformer(FastModel): ): modules_json_path = os.path.join(model_name, "modules.json") else: - modules_json_path = hf_hub_download(model_name, "modules.json", token = token) + modules_json_path = hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) with open(modules_json_path, "r", encoding = "utf-8") as f: modules_config = json.load(f) @@ -571,6 +583,8 @@ class FastSentenceTransformer(FastModel): model_name, os.path.join(pooling_path, "config.json"), token = token, + cache_dir = cache_dir, + revision = revision, ) break @@ -950,7 +964,12 @@ class FastSentenceTransformer(FastModel): f.write(content) @staticmethod - def _module_path(model_name, token = None): + def _module_path( + model_name, + token = None, + cache_dir = None, + revision = None, + ): """Return the path to the modules.json file, or None.""" try: if os.path.exists(model_name) and os.path.isdir(model_name): @@ -958,7 +977,13 @@ class FastSentenceTransformer(FastModel): return path if os.path.exists(path) else None else: try: - return hf_hub_download(model_name, "modules.json", token = token) + return hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) except: return None except: @@ -1135,6 +1160,8 @@ class FastSentenceTransformer(FastModel): max_seq_length, pooling_mode, trust_remote_code = False, + cache_dir = None, + revision = None, ) -> tuple[OrderedDict, bool]: """Load modules from modules.json, else fall back to hard-coded modules. @@ -1145,7 +1172,9 @@ class FastSentenceTransformer(FastModel): from sentence_transformers.models import Pooling, Normalize modules = OrderedDict() - modules_json_path = FastSentenceTransformer._module_path(model_name, token) + modules_json_path = FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) if modules_json_path: with open(modules_json_path, encoding = "utf8") as f: @@ -1171,7 +1200,13 @@ class FastSentenceTransformer(FastModel): load_path = os.path.join(model_name, module_path) else: try: - load_path = load_dir_path(model_name, module_path, token = token) + load_path = load_dir_path( + model_name, + module_path, + token = token, + cache_folder = cache_dir, + revision = revision, + ) except Exception as e: print(f"Unsloth Warning: Could not download module {module_path}: {e}") continue @@ -1198,7 +1233,9 @@ class FastSentenceTransformer(FastModel): hidden_size = getattr(model.config, "hidden_size", 768) if pooling_mode == "mean": - pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) + pooling_mode = FastSentenceTransformer._read_pooling_mode( + model_name, token, cache_dir = cache_dir, revision = revision + ) modules["1"] = Pooling(word_embedding_dimension = hidden_size, pooling_mode = pooling_mode) modules["2"] = Normalize() @@ -1386,6 +1423,45 @@ class FastSentenceTransformer(FastModel): "Run `pip install sentence-transformers` to install it." ) + # Validate the load modes BEFORE the prefetch so a bad config fails without downloading weights. + # Guard on not for_inference: that branch below never used these flags. + if not for_inference: + # sanity check, thanks Etherl: + if full_finetuning and (load_in_4bit or load_in_8bit): + print( + "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." + ) + load_in_4bit = False + load_in_8bit = False + load_in_fp8 = False + load_in_16bit = False + + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: + raise RuntimeError( + "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" + "Also, we by default set `load_in_16bit = True`.\n" + "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" + "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" + ) + + # Prefetch so the ST load below is a cache hit. weights_at_root stays False (ST component + # weights live in per-module subfolders). Resolve the same cache the load uses: HF cache_dir, + # else cache_folder, else SENTENCE_TRANSFORMERS_HOME, else default -- a wrong cache misses the warm. + _st_prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + local_files_only = kwargs.get("local_files_only", False), + # Forward force_download so the refresh happens in the killable child, then clear it so the + # in-process ST load reuses the warm cache instead of re-downloading over unguarded Xet. + force_download = kwargs.get("force_download", False), + ) + if _st_prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues if for_inference: st_device = device_map @@ -1416,27 +1492,16 @@ class FastSentenceTransformer(FastModel): if k in kwargs: st_kwargs[k] = kwargs[k] + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). + _st_cache = kwargs.get("cache_dir") or kwargs.get("cache_folder") + if _st_cache is not None: + st_kwargs["cache_folder"] = _st_cache + st_model = SentenceTransformer(model_name, **st_kwargs) return st_model - # sanity check, thanks Etherl: - if full_finetuning and (load_in_4bit or load_in_8bit): - print( - "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." - ) - load_in_4bit = False - load_in_8bit = False - load_in_fp8 = False - load_in_16bit = False - - if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: - raise RuntimeError( - "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" - "Also, we by default set `load_in_16bit = True`.\n" - "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" - "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" - ) - + # Load-mode validation already ran before the prefetch above. if "auto_model" not in kwargs: kwargs["auto_model"] = AutoModel @@ -1533,7 +1598,8 @@ class FastSentenceTransformer(FastModel): elif is_mpnet: FastSentenceTransformer._patch_mpnet_v5() - # Load via native SentenceTransformer (bypasses Unsloth patching) + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). st_model = SentenceTransformer( model_name, device = st_device, @@ -1541,6 +1607,7 @@ class FastSentenceTransformer(FastModel): token = token, revision = revision, model_kwargs = model_kwargs, + cache_folder = kwargs.get("cache_dir") or kwargs.get("cache_folder"), ) # Store metadata for get_peft_model @@ -1646,7 +1713,18 @@ class FastSentenceTransformer(FastModel): # No modules.json -> force 16-bit: saving is custom for these models and # 4-bit would need dequant in save_pretrained_merged, not worth it. - has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None + # Resolve the warmed cache: hf_hub_download ignores SENTENCE_TRANSFORMERS_HOME, so pass it as cache_dir. + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, + token, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + revision = revision, + ) + is not None + ) if not has_modules_json and load_in_4bit: print( @@ -1656,6 +1734,12 @@ class FastSentenceTransformer(FastModel): load_in_4bit = False load_in_16bit = True + # The fallback FastModel load reads HF cache_dir, not ST's cache_folder/SENTENCE_TRANSFORMERS_HOME. + # Point it at the warmed cache, but only when no explicit cache_dir was passed (which wins). + _st_cache_dir = kwargs.get("cache_folder") or os.environ.get("SENTENCE_TRANSFORMERS_HOME") + if _st_cache_dir is not None and "cache_dir" not in kwargs: + kwargs["cache_dir"] = _st_cache_dir + try: model, tokenizer = FastModel.from_pretrained( model_name = model_name, @@ -1697,6 +1781,12 @@ class FastSentenceTransformer(FastModel): max_seq_length, pooling_mode, trust_remote_code = trust_remote_code, + # Same resolved cache as above so the fallback module loads hit the warm, not Xet. + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + # Same revision as the weight load so modules hit the warm (None = default branch). + revision = revision, ) st_device = device_map diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 689e362f95..0a68a49fee 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -37,6 +37,7 @@ from ._utils import ( _get_text_only_config, _is_family_text_decoder, _apply_text_only_key_mapping, + _select_moe_detection_targets, set_task_config_attr, ) from ._utils import * @@ -551,6 +552,7 @@ def _construct_vlm_processor_fallback( model_type, token, trust_remote_code, + cache_dir = None, local_files_only = False, ): """Build a VLM processor manually when AutoProcessor.from_pretrained fails (some VLMs @@ -568,6 +570,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check) @@ -576,6 +579,7 @@ def _construct_vlm_processor_fallback( padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Read tokenizer_config.json for special tokens: prefer the local file (offline @@ -601,6 +605,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, "tokenizer_config.json", token = token, + cache_dir = cache_dir, local_files_only = local_files_only, ) with open(config_path, "r", encoding = "utf-8") as f: @@ -632,6 +637,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) proc_class_name = PROCESSOR_MAPPING_NAMES.get(config.model_type) @@ -872,6 +878,9 @@ class FastBaseModel: # For debugging - we use a download counter to see if environments are not breaking or if HF is down get_statistics(kwargs.get("local_files_only", False)) + # The base + tokenizer prefetch runs AFTER the load-mode validation below, so an invalid + # load_in_* combination fails without first downloading a snapshot. + if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 elif os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": @@ -968,6 +977,53 @@ class FastBaseModel: raise RuntimeError( "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!" ) + + # Prefetch the repo (killable child) so the in-process load below is a cache hit. vLLM owns the + # weight download only when actually available; if fast_inference was requested but vLLM is + # missing, the load falls through in-process, so weights must still be warmed here. + _vllm_owns_weights = fast_inference and is_vLLM_available() + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + fast_inference = _vllm_owns_weights, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Warm a SEPARATE tokenizer repo only (model_name is covered above). Not model_name here: this + # runs before fast_inference_setup may remap the repo, so it would warm the wrong one. + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and _tokenizer_repo != model_name + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + _skip_modules = SKIP_QUANTIZATION_MODULES.copy() # Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers. # Mamba fused kernels pass out_proj.weight directly to F.linear, @@ -1278,6 +1334,18 @@ class FastBaseModel: # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # On the vLLM path the tokenizer warm was deferred (fast_inference_setup may remap model_name). + # Warm the now-final tokenizer repo so the load below hits the cache (a cached/local repo is a no-op). + if _vllm_owns_weights and isinstance(tokenizer_name, str) and tokenizer_name: + maybe_prefetch_hf_snapshot( + tokenizer_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + # Fix _Unsloth_Patched_ prefix in local config files from old saves (issue #4085) if os.path.isdir(tokenizer_name): import json as _json @@ -1315,6 +1383,7 @@ class FastBaseModel: language = whisper_language, task = whisper_task, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1327,6 +1396,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1337,6 +1407,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1355,6 +1426,7 @@ class FastBaseModel: model_type_arch, token, trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _fe: @@ -1440,6 +1512,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = local_files_only, ) model, _fallback_tok = patch_tokenizer(model, _fallback_tok) @@ -1469,6 +1542,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1478,6 +1552,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) @@ -1629,6 +1704,16 @@ class FastBaseModel: ) else: _audio_kwargs = {} + # Remember the caller's ORIGINAL explicit leaf list for MoE expert + # detection. When an explicit list is routed through get_peft_regex for + # family scoping below, the generated regex carries get_peft_regex's full + # "mlp|feed_forward|ffn|dense" component block even when the caller named + # only attention leaves (q/k/v/o_proj). Keying expert detection on that + # regex would train the experts for an attention-only request. The + # original list carries the true leaf intent, so use it for MoE detection; + # only the auto (None / "all-linear") path relies on the regex, whose mlp + # block is the sole remaining MLP-intent signal on fused-expert models. + _moe_detect_target = target_modules if type(target_modules) in (list, tuple) else None if target_modules is None or target_modules == "all-linear": target_modules = get_peft_regex( model, @@ -1706,9 +1791,21 @@ class FastBaseModel: loftq_config, lora_dropout, bias, init_lora_weights, model ) - # Auto-detect MoE models and populate target_parameters for expert layers + # Auto-detect MoE models and populate target_parameters for expert layers. + # Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + # attention-only request does not train experts via get_peft_regex's mlp block, + # but only when MLP and language families are both still in scope. If the caller + # scoped MLP or language OFF (finetune_mlp_modules / finetune_language_layers + # False), the scoped regex already dropped the experts, so honor it instead of + # re-introducing the original list's gate/up/down leaves. if target_parameters is None: - target_parameters = get_moe_target_parameters(model, target_modules) + _moe_targets = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + target_parameters = get_moe_target_parameters(model, _moe_targets) if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model) @@ -2207,3 +2304,16 @@ def check_dataset_for_missing_videos( warnings.warn(error_msg, stacklevel = 2) return missing + + +# Auto-enable grouped-GEMM MoE (transformers<5 ModuleList experts); see llama.py. +try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe + FastBaseModel.from_pretrained = staticmethod( + wrap_loader_for_grouped_moe(FastBaseModel.from_pretrained) + ) + FastBaseModel.get_peft_model = staticmethod( + wrap_loader_for_grouped_moe(FastBaseModel.get_peft_model) + ) +except Exception: + pass diff --git a/unsloth/save.py b/unsloth/save.py index 50ae4119bd..020c63a9e2 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -223,6 +223,49 @@ def _normalize_torchao_method(save_method): return TORCHAO_EXPORT_SCHEMES.get(key) +def _loaded_via_remote_code(obj): + """True if `obj`'s class comes from downloaded custom code (an auto_map module). + + Transformers loads auto_map code into the ``transformers_modules`` package, so a + ``transformers_modules`` class proves the original load actually ran that remote code + (which the caller's / Studio's consent gate scans at load time). Export paths derive their + reload trust_remote_code from this - the already approved load decision - instead of from a + checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its + unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper + layers so a LoRA over a custom-code base is still detected, and processor components so a + custom tokenizer held inside a built-in processor keeps its approved trust. + """ + seen = set() + queue = [obj] + while queue and len(seen) < 16: + node = queue.pop(0) + if node is None or id(node) in seen: + continue + seen.add(id(node)) + # __module__ can be None/absent on some dynamically created or C-extension classes; + # treat anything non-string as "not remote code" rather than crashing the export. + module = getattr(type(node), "__module__", None) + if isinstance(module, str) and module.startswith("transformers_modules"): + return True + if hasattr(node, "get_base_model"): + try: + queue.append(node.get_base_model()) + except Exception: + pass + # PEFT / trainer wrappers hold the real model in base_model / model; a built-in + # ProcessorMixin holds its (possibly custom-code) components as attributes. + for attr in ( + "base_model", + "model", + "tokenizer", + "image_processor", + "feature_extractor", + "video_processor", + ): + queue.append(getattr(node, attr, None)) + return False + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -2883,15 +2926,16 @@ def unsloth_save_pretrained_gguf( "Unsloth: quantization_method can only be a string or a list of strings" ) for i, quant_method in enumerate(quantization_method): - quant_method = quant_method.lower() + if quant_method is None: + quant_method = "q8_0" + else: + quant_method = quant_method.lower() if quant_method == "not_quantized": quant_method = "f16" elif quant_method == "fast_quantized": quant_method = "q8_0" elif quant_method == "quantized": quant_method = "q4_k_m" - elif quant_method is None: - quant_method = "q8_0" quantization_methods.append(quant_method.lower()) try: @@ -3532,7 +3576,9 @@ def _unsloth_save_lora_gguf( cmd += ["--base", base_model_id] else: cmd += ["--base-model-id", base_model_id] - if bool(getattr(model.config, "auto_map", None)): + # Only pass --trust-remote-code when the loaded model actually came from custom code (the + # approved load decision), not merely because its config carries an auto_map entry. + if _loaded_via_remote_code(model): cmd.append("--trust-remote-code") # Expose the token to the converter so it can fetch a gated/private base config from the Hub. @@ -3682,15 +3728,16 @@ def save_to_gguf_generic( "Unsloth: quantization_method can only be a string or a list of strings" ) for i, quant_method in enumerate(quantization_method): - quant_method = quant_method.lower() + if quant_method is None: + quant_method = "q8_0" + else: + quant_method = quant_method.lower() if quant_method == "not_quantized": quant_method = "f16" elif quant_method == "fast_quantized": quant_method = "q8_0" elif quant_method == "quantized": quant_method = "q4_k_m" - elif quant_method is None: - quant_method = "q8_0" new_quantization_methods.append(quant_method.lower()) else: new_quantization_methods.append(quantization_type.lower()) @@ -4387,8 +4434,8 @@ def _unsloth_save_compressed_tensors( ) unsloth_generic_save(**merge_args) - # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal - # model exposes a vision_config or an explicitly vision-named architecture; a bare + # 4) Detect VLM from the in-memory model config. A vision/multimodal model exposes a + # vision_config or an explicitly vision-named architecture; a bare # *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it # is not treated as a VLM on its own. is_vlm = False @@ -4402,9 +4449,13 @@ def _unsloth_save_compressed_tensors( "Unsloth: FP8/FP4 compressed export for vision / multimodal models is " "experimental; vision-tower layers may be affected." ) - trust_remote_code = ( - bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False - ) + # trust_remote_code must reflect the approved load decision (whether the model / tokenizer + # was actually loaded from custom code), not the config's static auto_map, so a + # built-in-loadable model carrying auto_map cannot run unvetted code in the subprocess. + # Model and tokenizer trust stay separate, like the torchao path: an approved custom + # tokenizer must not enable an unapproved model's code in the subprocess (or vice versa). + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a # str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp. @@ -4479,8 +4530,10 @@ def _unsloth_save_compressed_tensors( cmd += ["--calibration-dataset", calib_value] if is_vlm: cmd.append("--is-vlm") - if trust_remote_code: + if model_trust: cmd.append("--trust-remote-code") + if tok_trust: + cmd.append("--trust-remote-code-tokenizer") if variant: cmd += ["--variant", variant] @@ -4679,35 +4732,19 @@ def _unsloth_save_torchao( ) unsloth_generic_save(**merge_args) - # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. - # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off - # vision_config / a vision-named architecture only, like the compressed path. + # 2) Detect VLM + reload class. A bare *ForConditionalGeneration also matches text seq2seq + # (T5/BART/Whisper), so key off vision_config / a vision-named architecture only. is_vlm = False - trust_remote_code = False if hasattr(model, "config"): archs = getattr(model.config, "architectures", None) or [] is_vlm = hasattr(model.config, "vision_config") or any( x.endswith("ForVisionText2Text") for x in archs ) - trust_remote_code = bool(getattr(model.config, "auto_map", None)) - # Custom code can be declared only in the tokenizer/processor config, so also honor an - # auto_map in any staged config (the original load already had the user's consent). - if not trust_remote_code: - for _cfg in ( - "config.json", - "tokenizer_config.json", - "processor_config.json", - "preprocessor_config.json", - ): - try: - _p = os.path.join(staging, _cfg) - if os.path.exists(_p): - with open(_p, "r", encoding = "utf-8") as _f: - if "auto_map" in json.load(_f): - trust_remote_code = True - break - except Exception: - pass + # trust_remote_code must reflect the approved load decision - whether the in-memory model / + # tokenizer was itself loaded from custom code - not the staged config's auto_map, which an + # attacker can set on a built-in-loadable model to run unvetted code past the consent gate. + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) # Reload with the class that matches the checkpoint: an image-text VLM class (with a # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and @@ -4766,12 +4803,10 @@ def _unsloth_save_torchao( staging, device_map = "auto", quantization_config = TorchAoConfig(quant_type = quant_type), - trust_remote_code = trust_remote_code, + trust_remote_code = model_trust, **dtype_kw, ) - staged_tokenizer = auto_processor.from_pretrained( - staging, trust_remote_code = trust_remote_code - ) + staged_tokenizer = auto_processor.from_pretrained(staging, trust_remote_code = tok_trust) quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) staged_tokenizer.save_pretrained(out_dir) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 93dfa9b2ad..3a91ef188d 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -563,8 +563,11 @@ def _load_correct_tokenizer( # /tmp of Kaggle seems has a 80GB limit! # Let's utilize them cache_dir = os.path.join(KAGGLE_TMP, cache_dir) - else: + elif cache_dir == "huggingface_tokenizers_cache": + # This default name is Colab/Kaggle-only; elsewhere use the HF default cache. cache_dir = None + # else: keep a caller-supplied cache_dir so the tokenizer loads from the prefetch-warmed dir instead + # of risking an in-process Hub/Xet transfer. # Try loading the slow tokenizer. If it fails, then try Fast only # Mainly to solve Deepseek models with no tokenizer.model file @@ -1323,6 +1326,7 @@ def check_tokenizer( padding_side = "right", token = None, _reload = True, + cache_dir = None, ): # Checks tokenizer for out of bounds ids. # Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha @@ -1413,10 +1417,11 @@ def check_tokenizer( f"Fix your tokenizer since it'll perform out of bounds memory accesses." ) - if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT: - cache_dir = "huggingface_tokenizers_cache" - else: - cache_dir = None + # Reuse a caller-supplied cache_dir (warmed cache) for the repair reload; else the + # Colab/Kaggle sentinel (HF default elsewhere), as load_correct_tokenizer does. + reload_cache_dir = cache_dir + if reload_cache_dir is None and (IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT): + reload_cache_dir = "huggingface_tokenizers_cache" # Sometimes slow tokenizer does not work like Deepseek try: @@ -1430,7 +1435,7 @@ def check_tokenizer( use_fast = False, legacy = False, from_slow = True, - cache_dir = cache_dir, + cache_dir = reload_cache_dir, ) return check_tokenizer( model = model, @@ -1440,6 +1445,7 @@ def check_tokenizer( padding_side = padding_side, token = token, _reload = False, + cache_dir = cache_dir, ) break except: diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 2e984bad0a..68fb33dad9 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -17,6 +17,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Optional, Tuple @@ -42,6 +43,17 @@ if HAS_XFORMERS and torch.cuda.is_available(): HAS_XFORMERS = False SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") +# PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load +# torch flex_attention. +_flex_shared_prefix_attention = None +if os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ("0", "false", "no", "off"): + try: + from .prefix_grouper_kernel import ( + flex_shared_prefix_attention as _flex_shared_prefix_attention, + ) + except Exception: + _flex_shared_prefix_attention = None + FLASH_VARLEN = "flash_varlen" FLASH_DENSE = "flash_dense" XFORMERS = "xformers" @@ -84,6 +96,9 @@ class AttentionContext: attention_mask: Optional[Tensor] causal_mask: Optional[Any] sliding_window: Optional[int] = None + # PrefixGrouper: non-None routes Q/K/V through the FlexAttention shared-prefix kernel; + # None leaves every existing construction/behavior unchanged. + prefix_seg_info: Optional[Any] = None def select_attention_backend(use_varlen: bool = False) -> str: @@ -99,6 +114,33 @@ def select_attention_backend(use_varlen: bool = False) -> str: return SDPA +def resolve_prefix_seg_info(kwargs, past_key_value, attention_mask): + """PrefixGrouper shared-prefix segment table resolver for the arch attention forwards. + + The GRPO PrefixGrouper packed path rides a ``PrefixSegInfo`` in through ``**kwargs`` + (same route as ``packed_seq_lengths``). When present, the forward must route Q/K/V + through the FlexAttention shared-prefix kernel via ``AttentionContext.prefix_seg_info``. + + Returns the seg table (or ``None`` when PrefixGrouper did not group this batch -- the + unchanged path). Hardened: the shared-prefix stream is NOT a plain causal sequence, so running + it under a KV cache or an explicit padding mask would silently produce wrong logprobs. + That combination can only arise from misuse (PrefixGrouper only rides in via the GRPO + logprob forward, which is mask-free prefill), so we RAISE loudly instead of degrading + to a wrong result. + + Factored here so every arch (llama/mistral/qwen3/gemma2/cohere/granite/falcon_h1) + shares one implementation and cannot drift. + """ + seg = kwargs.get("prefix_seg_info", None) + if seg is not None and (past_key_value is not None or attention_mask is not None): + raise RuntimeError( + "PrefixGrouper: prefix_seg_info requires prefill with no KV cache and no " + f"attention_mask (got past_key_value={past_key_value is not None}, " + f"attention_mask={attention_mask is not None})." + ) + return seg + + def run_attention( *, config: AttentionConfig, context: AttentionContext, Q: Tensor, K: Tensor, V: Tensor ) -> Tensor: @@ -111,6 +153,28 @@ def run_attention( and SDPA handle packing via a block-diagonal mask. """ + # PrefixGrouper shared-prefix attention (GRPO dedup). Q/K/V here are [bsz, H, T, D]; + # the kernel takes/returns [1, T, H, D], matching the other backends. The field is + # only set when the env gate is on and grouping succeeded; None keeps every backend + # byte-identical. + if context.prefix_seg_info is not None: + flex_shared_prefix_attention = _flex_shared_prefix_attention + if flex_shared_prefix_attention is None: + # gate flipped on after import (or one-time load failed): resolve lazily. + from ..utils.prefix_grouper_kernel import flex_shared_prefix_attention + + scale = None + if config.flash_varlen_kwargs: + scale = config.flash_varlen_kwargs.get("softmax_scale") + A = flex_shared_prefix_attention( + Q.transpose(1, 2), + K.transpose(1, 2), + V.transpose(1, 2), + context.prefix_seg_info, + scale = scale, + ) + return A # [1, T, n_heads, head_dim] + backend = config.backend if backend == FLASH_VARLEN and context.seq_info is None: backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA @@ -337,5 +401,6 @@ __all__ = [ "AttentionConfig", "AttentionContext", "select_attention_backend", + "resolve_prefix_seg_info", "run_attention", ] diff --git a/unsloth/utils/prefix_grouper.py b/unsloth/utils/prefix_grouper.py new file mode 100644 index 0000000000..4e6ff9672c --- /dev/null +++ b/unsloth/utils/prefix_grouper.py @@ -0,0 +1,351 @@ +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""PrefixGrouper layout builder + completion-logprob extraction for the Unsloth GRPO +packed path (all archs that route through the varlen attention dispatch). + +Given the de-padded, LEFT-PACKED input_ids the packed GRPO path already works with, this +module: + + 1. Detects consecutive ``num_generations`` rows that share a prompt prefix (byte- + identical prompt precondition; falls back / returns None otherwise). + 2. Builds ONE flat shared-prefix stream across all groups + ``[ prefix_g0, suf_g0_0 .. suf_g0_{G-1}, prefix_g1, ... ]`` with position_ids that + continue each prefix positionally, plus a ``PrefixSegInfo`` segment table for the + FlexAttention shared-prefix kernel. + 3. Extracts completion logprobs via the index map (completion pos ``j==0`` predicted + from the shared prefix's last token; ``j>=1`` from the preceding suffix token) and + scatters them back into ``[total_rows, W]`` EXACTLY where the full-row packed path + puts them (dest = ``orig_row*L + orig_col``), so grpo_compute_loss / completion_mask + / TIS / metrics are byte-untouched. + +The flat stream is built by GATHERING original (row, col) coordinates out of input_ids, +so the grad path's autograd flows to the same embedding rows as today (the shared prefix +now contributes grad once = the sum of the G repeats, which is mathematically identical). + +``chunked_hidden_states_selective_log_softmax`` (from unsloth_zoo, passed in) is reused +verbatim over the gathered predicting-position hidden states, so fp32 accumulation, +logit_scale/softcapping/temperature are all preserved. + +Env: + UNSLOTH_GRPO_PREFIX_GROUPER=1 engage (default ON; set 0 to disable). Auto-off under vLLM. + UNSLOTH_GRPO_PREFIX_GROUPER_TOKR=1.3 tok_r auto-gate threshold (env-overridable) + UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY=1 first-step self-verify (default ON) + UNSLOTH_GRPO_PREFIX_GROUPER_TOL=0.7 self-verify PASS band (nats) +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch + +from .prefix_grouper_kernel import build_seg_info_multigroup, PrefixSegInfo + + +# --------------------------------------------------------------------------- +# Env helpers +# --------------------------------------------------------------------------- +def env_on(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() not in ("0", "false", "no", "off") + + +# One-time env reads; the helpers stay callable since unsloth_zoo imports and calls them. +_ENABLED = env_on("UNSLOTH_GRPO_SEQ_PACKING", "1") and env_on("UNSLOTH_GRPO_PREFIX_GROUPER", "1") +_VERIFY_ON = env_on("UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY", "1") +_TOKR_THRESHOLD = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOKR", "1.3")) +_TOL_OK = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOL", "0.7")) + + +def prefix_grouper_enabled() -> bool: + """PrefixGrouper requires seq-packing on (it reuses its de-pad + scatter machinery).""" + return _ENABLED + + +def verify_on() -> bool: + return _VERIFY_ON + + +def tokr_threshold() -> float: + return _TOKR_THRESHOLD + + +def tol_ok() -> float: + return _TOL_OK + + +# diff >= TOL_KILL = broken mask/isolation -> structure permanently unsafe; between +# tol_ok and TOL_KILL -> fall back for this shape but keep trying others. +TOL_KILL = 1.5 + + +@dataclass +class GroupLayout: + """Everything the GRPO forward needs to run + extract the shared-prefix path.""" + + flat_ids: torch.Tensor # [1, T] (T == seg.T) + position_ids: torch.Tensor # [1, T] + prefix_seg_info: PrefixSegInfo + # per completion target token, aligned 1:1: + tgt_rows: torch.Tensor # [N] original row index + tgt_cols: torch.Tensor # [N] original padded column in that row + tgt_pred: torch.Tensor # [N] flat predicting index (into the T stream) + tgt_flat: torch.Tensor # [N] flat index of the target token itself (into T) + total_rows: int + L: int # original padded seq length (input_ids.shape[1]) + W: int # logits_to_keep + max_left_pad (scatter width) + tok_r: float + signature: Tuple + + def extract_logps( + self, + hidden, + lm_head, + chunked_fn, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) -> torch.Tensor: + """hidden: [1, T, Hdim] (pre-lm_head hidden states, UNSLOTH_RETURN_HIDDEN_STATES=1). + Returns [total_rows, W] float32, byte-compatible with the packed path result.""" + # In a sharded model hidden may live on the lm-head device; move the small index + # maps to hidden.device before indexing. + device = hidden.device + pred_h = hidden[0, self.tgt_pred.to(device), :].unsqueeze(0) # [1, N, Hdim] + tgt_ids = self.flat_ids[0, self.tgt_flat].to(device).unsqueeze(0) # [1, N] + sel = chunked_fn( + pred_h, + lm_head, + tgt_ids, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] # [N] logprobs + dest = self.tgt_rows.to(device) * self.L + self.tgt_cols.to(device) + result = ( + torch.zeros(self.total_rows * self.L, dtype = torch.float32, device = device) + .index_put((dest,), sel.to(torch.float32)) + .view(self.total_rows, self.L)[:, -self.W :] + ) + return result + + +def _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows): + """CPU-side grouping. Returns group dicts or None. Mirrors the packed _pk_* partition. + + A row's REAL tokens are the columns where input != pad. Its completion region (what + the packed path scatters, then completion_mask masks) is the real columns with + original col >= cstart_r, where cstart_r = (L - logits_to_keep) - left_pad_r. The + prompt is the real columns < cstart_r. Within a GRPO group all G rows share the same + prompt => same left_pad => same cstart => the prompt real columns are BYTE-IDENTICAL + across the group (the shared prefix). We require that byte-identity (falls back + otherwise). No prompt-tail special-casing: every suffix token is scattered exactly + like the packed path; completion_mask masks the leading prompt-tail positions. + """ + G = num_generations + if G is None or G < 2 or total_rows % G != 0: + return None + groups = [] + for g0 in range(0, total_rows, G): + rows = list(range(g0, g0 + G)) + prompt_cols_per_row = [] # real cols < cstart + prompt_toks_per_row = [] + comp_cols_per_row = [] # real cols >= cstart (the completion region packed scatters) + for r in rows: + cs = cstart_cpu[r] + rc = real_cols_cpu[r] + p_cols = [c for c in rc if c < cs] + c_cols = [c for c in rc if c >= cs] + prompt_cols_per_row.append(p_cols) + prompt_toks_per_row.append([ids_cpu[r][c] for c in p_cols]) + comp_cols_per_row.append(c_cols) + if any(len(p) == 0 for p in prompt_toks_per_row): + return None + # require BYTE-IDENTICAL prompts across the group (shared-prefix precondition). + P = len(prompt_toks_per_row[0]) + if any(len(prompt_toks_per_row[k]) != P for k in range(1, G)): + return None + p0 = prompt_toks_per_row[0] + if any(prompt_toks_per_row[k] != p0 for k in range(1, G)): + return None + if P == 0: + return None + R_list = [len(c) for c in comp_cols_per_row] + if sum(R_list) == 0: + return None + groups.append( + dict( + rows = rows, + P = P, + prefix_cols = prompt_cols_per_row[0], # shared prompt real columns (row0) + prefix_row = rows[0], + R_list = R_list, + suf_cols = comp_cols_per_row, # per-row completion-region real columns + ) + ) + return groups + + +def _tok_r(groups) -> float: + tok_full = 0 + tok_sp = 0 + for gm in groups: + P = gm["P"] + Rs = gm["R_list"] + tok_full += sum(P + r for r in Rs) # G*P + sumR + tok_sp += P + sum(Rs) # P + sumR + return (tok_full / tok_sp) if tok_sp else 1.0 + + +def build_group_layout( + input_ids, + logits_to_keep, + pad_id, + num_generations, + left_pad_tokens_per_prompt, + *, + apply_tokr_gate = True, + max_segment_cap = None, +): + """Build the shared-prefix GroupLayout, or return None to fall back to the packed path. + + input_ids : [B, L]. GRPO's layout is left-padded in the prompt and right-padded in + the completion. Real tokens of a row are a contiguous run not necessarily + starting at column 0. + logits_to_keep : int + left_pad_tokens_per_prompt : [B] long tensor (per-row left-pad count in the prompt). + """ + device = input_ids.device + total_rows, L = input_ids.shape + keep = input_ids != pad_id + # completion start column per row (matches create_completion_attention_mask / _pk_cstart). + cstart = ((L - logits_to_keep) - left_pad_tokens_per_prompt).to(torch.long) + cstart_cpu = cstart.tolist() + ids_cpu = input_ids.tolist() + # per-row real (non-pad) columns. GRPO rows are one contiguous real run, so derive + # [first, first+n) on GPU; the O(B*L) scan is only a non-contiguous fallback. + n_real = keep.sum(dim = 1) + first = torch.argmax(keep.to(torch.int8), dim = 1) + ar = torch.arange(L, device = device) + contiguous = bool( + (keep == ((ar >= first.unsqueeze(1)) & (ar < (first + n_real).unsqueeze(1)))).all() + ) + if contiguous: + real_cols_cpu = [list(range(f, f + n)) for f, n in zip(first.tolist(), n_real.tolist())] + else: + keep_cpu = keep.tolist() + real_cols_cpu = [[c for c in range(L) if keep_cpu[r][c]] for r in range(total_rows)] + + groups = _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows) + if groups is None: + return None + + # sliding-window guard: a group's PG span is P + max(R); fall back if it exceeds the window. + if max_segment_cap is not None: + for gm in groups: + if gm["P"] + max(gm["R_list"]) > max_segment_cap: + return None + + tok_r = _tok_r(groups) + if apply_tokr_gate and tok_r < tokr_threshold(): + return None # low reuse -> not worth it; use the full-row packed path + + # Build flat stream by gathering original (row, col) coordinates. + group_specs = [(gm["P"], gm["R_list"]) for gm in groups] + seg, group_meta = build_seg_info_multigroup(group_specs, device) + + flat_src_rows: List[int] = [] + flat_src_cols: List[int] = [] + pos_list: List[int] = [] + tgt_rows: List[int] = [] + tgt_cols: List[int] = [] + tgt_pred: List[int] = [] + tgt_flat: List[int] = [] + + for gm, meta in zip(groups, group_meta): + rows = gm["rows"] + P = gm["P"] + r0 = gm["prefix_row"] + prefix_cols = gm["prefix_cols"] # ORIGINAL real prompt columns (len P) of row0 + plast = meta["prefix_last_index"] # base + P - 1 + # gather the shared prefix once, from row0. + flat_src_rows.extend([r0] * P) + flat_src_cols.extend(prefix_cols) + pos_list.extend(range(P)) + # suffixes: every suffix token is a completion-region target (scattered like the + # packed path; completion_mask hides prompt-tail positions). + for i, r in enumerate(rows): + cols = gm["suf_cols"][i] + r_i = len(cols) + s, e = meta["suffix_slices"][i] # flat offsets [s, e) + flat_src_rows.extend([r] * r_i) + flat_src_cols.extend(cols) + pos_list.extend(range(P, P + r_i)) + for j in range(r_i): + # pos 0 is predicted from the prefix's last token; j>=1 from the previous suffix token. + pred = plast if j == 0 else (s + j - 1) + tgt_rows.append(r) + tgt_cols.append(cols[j]) # ORIGINAL padded column in row r + tgt_pred.append(pred) + tgt_flat.append(s + j) # flat index of the target token itself + + T = len(flat_src_rows) + assert T == seg.T, f"flat stream len {T} != seg.T {seg.T}" + fr = torch.tensor(flat_src_rows, device = device, dtype = torch.long) + fc = torch.tensor(flat_src_cols, device = device, dtype = torch.long) + flat_ids = input_ids[fr, fc].unsqueeze(0) # [1, T] (grad-safe gather) + position_ids = torch.tensor(pos_list, device = device, dtype = torch.long).unsqueeze(0) + + max_left_pad = int(left_pad_tokens_per_prompt.max().item()) if total_rows else 0 + W = logits_to_keep + max_left_pad + + # self-verify cache key: the mask/index-map/scatter logic is structural, so key on + # (num_groups, group_sizes), not exact lengths -- GRPO lengths change every step and + # keying on T would re-verify forever ("verify once, then trust", like the packed path). + grp_sizes = tuple(sorted(len(gm["R_list"]) for gm in groups)) + sig = (len(groups), grp_sizes) + + return GroupLayout( + flat_ids = flat_ids, + position_ids = position_ids, + prefix_seg_info = seg, + tgt_rows = torch.tensor(tgt_rows, device = device, dtype = torch.long), + tgt_cols = torch.tensor(tgt_cols, device = device, dtype = torch.long), + tgt_pred = torch.tensor(tgt_pred, device = device, dtype = torch.long), + tgt_flat = torch.tensor(tgt_flat, device = device, dtype = torch.long), + total_rows = total_rows, + L = L, + W = W, + tok_r = tok_r, + signature = sig, + ) + + +__all__ = [ + "GroupLayout", + "build_group_layout", + "prefix_grouper_enabled", + "verify_on", + "tokr_threshold", + "tol_ok", + "TOL_KILL", + "env_on", +] diff --git a/unsloth/utils/prefix_grouper_kernel.py b/unsloth/utils/prefix_grouper_kernel.py new file mode 100644 index 0000000000..9a9719b015 --- /dev/null +++ b/unsloth/utils/prefix_grouper_kernel.py @@ -0,0 +1,436 @@ +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""FlexAttention shared-prefix kernel for PrefixGrouper (GRPO shared-prompt dedup). + +In GRPO every prompt spawns ``G = num_generations`` completions that share the same +prompt prefix. The full-row packed path forwards the identical prefix ``G`` times. +PrefixGrouper stores the prefix ONCE and concatenates only the ``G`` suffixes, with an +attention layout where each suffix token attends to ``[the single shared prefix] + +[causal within its own suffix]``. This kernel expresses that one-prefix -> many-suffix +fan-out via a ``torch.nn.attention.flex_attention`` block mask, so the masked-out +cross-suffix / cross-group blocks are never computed and the ``P + G*R`` FLOP saving is +realised (not merely a masked dense ``O(T^2)``). + +Mask semantics (identical to the certified SDPA oracle): + + keep(q_idx, kv_idx) = same_group(q, kv) AND + ( is_prefix[kv_idx] # full prefix visibility + OR ( suffix_of_kv[kv_idx] == suffix_of_kv[q_idx] # same suffix ... + AND kv_idx <= q_idx ) ) # ... causal within it + +This module is self-contained (no dependency on any temp/ scratch dir) so PrefixGrouper +works from the installed source after a fresh compile. It is only imported lazily from +``attention_dispatch.run_attention`` when ``prefix_seg_info`` is present, which itself is +only ever set when ``UNSLOTH_GRPO_PREFIX_GROUPER`` is on and grouping succeeded, so the +default (off) path never touches this file. + +Provided entry points: + * ``PrefixSegInfo`` : per-flat-token segment metadata + cache signature. + * ``build_seg_info_multigroup``: build PrefixSegInfo for many groups packed flat. + * ``build_seg_info_from_layout``: build PrefixSegInfo for ONE group (test helper). + * ``get_block_mask`` : cached create_block_mask keyed on the signature. + * ``flex_shared_prefix_attention(Q, K, V, prefix_seg_info)`` + Q/K/V of shape [1, T, n_heads, head_dim]; returns [1, T, n_heads, head_dim], + IDENTICAL semantics to the SDPA oracle. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch +from torch.nn.attention.flex_attention import ( + BlockMask, + create_block_mask, + flex_attention, +) + +# GRPO feeds many distinct segment lengths; at dynamo's default recompile_limit (8) the +# compiled kernel silently reuses a mismatched specialisation (wrong results). Raise it. +torch._dynamo.config.recompile_limit = max(getattr(torch._dynamo.config, "recompile_limit", 8), 256) +torch._dynamo.config.accumulated_recompile_limit = max( + getattr(torch._dynamo.config, "accumulated_recompile_limit", 256), 2048 +) + + +# Compiled kernels: torch.compile fuses the sparse mask into one kernel. dynamic=True is +# required: T changes almost every GRPO batch and dynamic=False recompiles per T (~14s +# each). T is still padded to a multiple of 128 (_pad_len) for the backward kernel. +_flex_attention_compiled = torch.compile(flex_attention, dynamic = True) +_create_block_mask_compiled = torch.compile(create_block_mask, dynamic = True) + +# Flash block sizes by Q dtype (env-overridable). The two disjoint key runs (prefix + +# own-suffix) stress online-softmax accumulation: fp32 needs 32/32 for a ~1e-6 floor; +# bf16 passes parity at 128/64 and is ~5x faster (128/128 OOMs Triton on B200). +_FP32_BLOCK_M = int(os.environ.get("PG_FLEX_BLOCK_M", "32")) +_FP32_BLOCK_N = int(os.environ.get("PG_FLEX_BLOCK_N", "32")) +_BF16_BLOCK_M = int(os.environ.get("PG_FLEX_BF16_BLOCK_M", "128")) +_BF16_BLOCK_N = int(os.environ.get("PG_FLEX_BF16_BLOCK_N", "64")) + + +def _kernel_options_for_dtype(dtype): + """Pick the numerically-safe flash block sizes for the Q dtype.""" + if dtype == torch.bfloat16 or dtype == torch.float16: + return {"BLOCK_M": _BF16_BLOCK_M, "BLOCK_N": _BF16_BLOCK_N} + return {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} + + +# Backward-compat constant (fp32 default). +_FLEX_KERNEL_OPTIONS = {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} + +# The compiled backward trips an Inductor assertion when T is not a multiple of 128, so +# pad the flat sequence. Pad tokens form a group that attends to / is attended by nothing +# (all-masked rows return 0, not NaN) and are sliced off the output. +_PAD_MULTIPLE = 128 +_PAD_GROUP = -99 # sentinel group id / suffix id for pad tokens + + +def _pad_len(T: int) -> int: + return ((T + _PAD_MULTIPLE - 1) // _PAD_MULTIPLE) * _PAD_MULTIPLE + + +# --------------------------------------------------------------------------- +# Segment metadata +# --------------------------------------------------------------------------- + + +@dataclass +class PrefixSegInfo: + """Per-flat-token segment metadata driving the shared-prefix block mask. + + The label tensors are 1-D of length ``T_pad`` (>= real ``T``, padded up to a multiple + of 128 so the backward kernel compiles). Positions ``[T:T_pad)`` are pad tokens + (group/suffix == _PAD_GROUP) that attend to nothing. + + Attributes + ---------- + group_of_kv : LongTensor [T_pad] + Group id per flat token (0..num_groups-1); _PAD_GROUP for pad tokens. + is_prefix : BoolTensor [T_pad] + True iff the token is a prefix token of its group (False for pad). + suffix_of_kv : LongTensor [T_pad] + Suffix id per flat token; -1 for prefix, _PAD_GROUP for pad. Suffix ids are + globally unique across groups. + signature : hashable + Cache key for the block mask (depends only on the labels + T_pad). + T : int + Real flat sequence length (Q/K/V of this length are padded internally). + T_pad : int + Padded length (multiple of 128) at which the block mask is built. + """ + + group_of_kv: torch.Tensor + is_prefix: torch.Tensor + suffix_of_kv: torch.Tensor + signature: Tuple + T: int + T_pad: int + + +def _pad_labels(group_of_kv, is_prefix, suffix_of_kv, device): + """Pad the label tensors up to a multiple of 128 with pad-token sentinels.""" + T = int(group_of_kv.numel()) + T_pad = _pad_len(T) + if T_pad == T: + return group_of_kv, is_prefix, suffix_of_kv, T, T_pad + pad = T_pad - T + group_of_kv = torch.cat( + [group_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] + ) + is_prefix = torch.cat([is_prefix, torch.zeros(pad, dtype = torch.bool, device = device)]) + suffix_of_kv = torch.cat( + [suffix_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] + ) + return group_of_kv, is_prefix, suffix_of_kv, T, T_pad + + +def build_seg_info_from_layout(layout, device: Optional[torch.device] = None) -> PrefixSegInfo: + """Build PrefixSegInfo for ONE group from an object with ``.flat_ids``, ``.P`` and + ``.suffix_slices`` (used by the parity test / oracle helpers).""" + if device is None: + device = layout.flat_ids.device + T = int(layout.flat_ids.shape[1]) + P = int(layout.P) + + group_of_kv = torch.zeros(T, dtype = torch.long, device = device) # single group -> 0 + is_prefix = torch.zeros(T, dtype = torch.bool, device = device) + is_prefix[:P] = True + suffix_of_kv = torch.full((T,), -1, dtype = torch.long, device = device) + for i, (s, e) in enumerate(layout.suffix_slices): + suffix_of_kv[s:e] = i + + group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( + group_of_kv, is_prefix, suffix_of_kv, device + ) + sig = ("single", T_pad, P, tuple((s, e) for (s, e) in layout.suffix_slices)) + return PrefixSegInfo( + group_of_kv = group_of_kv, + is_prefix = is_prefix, + suffix_of_kv = suffix_of_kv, + signature = sig, + T = T, + T_pad = T_pad, + ) + + +def build_seg_info_multigroup( + group_specs: List[Tuple[int, List[int]]], device: torch.device +) -> Tuple[PrefixSegInfo, List[dict]]: + """Build PrefixSegInfo for several shared-prefix groups packed block-diagonally. + + Parameters + ---------- + group_specs : list of (P_g, [R_{g,0}, R_{g,1}, ...]) + For each group: prefix length and the list of suffix lengths. + + Returns + ------- + seg : PrefixSegInfo + group_meta : list of dicts with 'base', 'P', 'prefix_last_index', 'suffix_slices' + (flat offsets), enough to build the completion index map. + """ + group_of_list = [] + is_prefix_list = [] + suffix_of_list = [] + group_meta = [] + + base = 0 + suffix_counter = 0 + sig_parts = [] + for gid, (P, R_list) in enumerate(group_specs): + # prefix + group_of_list.append(torch.full((P,), gid, dtype = torch.long, device = device)) + is_prefix_list.append(torch.ones(P, dtype = torch.bool, device = device)) + suffix_of_list.append(torch.full((P,), -1, dtype = torch.long, device = device)) + prefix_last_index = base + P - 1 + suffix_slices = [] + cursor = base + P + for r in R_list: + group_of_list.append(torch.full((r,), gid, dtype = torch.long, device = device)) + is_prefix_list.append(torch.zeros(r, dtype = torch.bool, device = device)) + suffix_of_list.append(torch.full((r,), suffix_counter, dtype = torch.long, device = device)) + suffix_slices.append((cursor, cursor + r)) + cursor += r + suffix_counter += 1 + group_meta.append( + { + "base": base, + "P": P, + "prefix_last_index": prefix_last_index, + "suffix_slices": suffix_slices, + } + ) + sig_parts.append((P, tuple(R_list))) + base = cursor + + group_of_kv = torch.cat(group_of_list) + is_prefix = torch.cat(is_prefix_list) + suffix_of_kv = torch.cat(suffix_of_list) + group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( + group_of_kv, is_prefix, suffix_of_kv, device + ) + sig = ("multi", T_pad, tuple(sig_parts)) + seg = PrefixSegInfo( + group_of_kv = group_of_kv, + is_prefix = is_prefix, + suffix_of_kv = suffix_of_kv, + signature = sig, + T = T, + T_pad = T_pad, + ) + return seg, group_meta + + +# --------------------------------------------------------------------------- +# Block-mask builder + cache, keyed on (signature, device): the mask depends only on the +# per-token labels and T, so it is reused across layers and steps. + +_BLOCK_MASK_CACHE: Dict[Tuple, BlockMask] = {} + + +def _make_mask_mod(group_of_kv, is_prefix, suffix_of_kv): + """Return a mask_mod closure over the (device) label tensors. + + keep(q, kv) = same_group AND + ( is_prefix[kv] AND kv <= q # causal within/ into prefix + OR ( suffix_of_kv[kv] == suffix_of_kv[q] # same suffix ... + AND (not is_prefix[q]) # q is a suffix token ... + AND kv <= q ) ) # ... causal within it + + The single ``kv <= q`` guard on the is_prefix branch gives BOTH prefix-causal + behaviour (a prefix q sees only earlier prefix tokens) AND full-prefix-visibility for + suffixes (every prefix index < every suffix index in a group, so kv <= q always holds + for a suffix q vs a prefix kv of its group), matching the SDPA oracle exactly. + """ + + def mask_mod(b, h, q_idx, kv_idx): + same_group = group_of_kv[q_idx] == group_of_kv[kv_idx] + kv_is_prefix = is_prefix[kv_idx] + causal = kv_idx <= q_idx + same_suffix = (suffix_of_kv[kv_idx] == suffix_of_kv[q_idx]) & (~is_prefix[q_idx]) + keep = same_group & ((kv_is_prefix & causal) | (same_suffix & causal)) + return keep + + return mask_mod + + +def get_block_mask( + seg: PrefixSegInfo, + device: torch.device, + compile_mask: bool = True, +) -> BlockMask: + """Return a cached BlockMask for the segment signature (built once, reused). + + CRITICAL: the block mask is cached and shared across BOTH the no-grad old/ref logprob + forward (which runs under torch.inference_mode) and the grad training forward. If the + mask were first built under inference_mode, its tensors would be INFERENCE tensors that + "cannot be saved for backward" when reused in the grad forward. We therefore build the + mask with inference mode explicitly DISABLED, so the same cached BlockMask is a normal + tensor usable by autograd. (The mask depends only on integer labels; it needs no grad.) + """ + key = (seg.signature, str(device)) + bm = _BLOCK_MASK_CACHE.get(key) + if bm is not None: + return bm + + # Move labels to the consumer (Q) device: with a sharded model the seg tensors live on + # input_ids.device and would index cross-device. Copies once per (signature, device). + # These copies must also run with inference mode DISABLED (same reason as the mask build): + # when this entry is first built under the no-grad old/ref forward's inference_mode and + # device != seg.device, a .to(device) copy would be an inference tensor that mask_mod + # captures, which then cannot be saved for backward when the grad training forward reuses + # the cached mask. + builder = _create_block_mask_compiled if compile_mask else create_block_mask + with torch.inference_mode(False): + mask_mod = _make_mask_mod( + seg.group_of_kv.to(device), seg.is_prefix.to(device), seg.suffix_of_kv.to(device) + ) + bm = builder( + mask_mod, + B = 1, + H = None, + Q_LEN = seg.T_pad, + KV_LEN = seg.T_pad, + device = device, + ) + # FIFO bound: GRPO lengths change nearly every step, so evict the oldest to cap GPU pins. + if len(_BLOCK_MASK_CACHE) >= 8: + _BLOCK_MASK_CACHE.pop(next(iter(_BLOCK_MASK_CACHE))) + _BLOCK_MASK_CACHE[key] = bm + return bm + + +def clear_block_mask_cache(): + _BLOCK_MASK_CACHE.clear() + + +def _pad_qkv_seq(x: torch.Tensor, T_pad: int) -> torch.Tensor: + """Zero-pad a [B, H, T, D] tensor along the sequence dim up to T_pad.""" + T = x.shape[2] + if T_pad == T: + return x + pad = torch.zeros(x.shape[0], x.shape[1], T_pad - T, x.shape[3], device = x.device, dtype = x.dtype) + return torch.cat([x, pad], dim = 2) + + +def _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad): + """Pad q/k/v to T_pad, run flex, slice the output back to T. q/k/v: [B,H,T,D].""" + qp = _pad_qkv_seq(q, T_pad) + kp = _pad_qkv_seq(k, T_pad) + vp = _pad_qkv_seq(v, T_pad) + if compiled: + out = _flex_attention_compiled( + qp, + kp, + vp, + block_mask = block_mask, + enable_gqa = enable_gqa, + scale = scale, + kernel_options = _kernel_options_for_dtype(qp.dtype), + ) + else: + # eager path (fp64 parity): dense scores, no kernel_options. + out = flex_attention( + qp, + kp, + vp, + block_mask = block_mask, + enable_gqa = enable_gqa, + scale = scale, + ) + return out[:, :, :T, :] + + +# --------------------------------------------------------------------------- +# The kernel entry point +# --------------------------------------------------------------------------- + + +def flex_shared_prefix_attention( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + prefix_seg_info: PrefixSegInfo, + scale: Optional[float] = None, + block_mask: Optional[BlockMask] = None, + compiled: bool = True, +) -> torch.Tensor: + """Shared-prefix attention via FlexAttention. + + Parameters + ---------- + Q, K, V : Tensor [1, T, n_heads, head_dim] + (Q has n_heads, K/V have n_kv_heads for GQA). + prefix_seg_info : PrefixSegInfo + scale : optional float, softmax scale (defaults to 1/sqrt(head_dim)). + block_mask : optional precomputed BlockMask (else built/cached from seg info). + + Returns + ------- + Tensor [1, T, n_heads, head_dim], identical semantics to the SDPA oracle branch. + """ + assert Q.dim() == 4 and Q.shape[0] == 1, f"expected [1,T,H,D], got {tuple(Q.shape)}" + device = Q.device + # FlexAttention wants [B, H, T, D]. + q = Q.transpose(1, 2) # [1, n_heads, T, D] + k = K.transpose(1, 2) # [1, n_kv_heads, T, D] + v = V.transpose(1, 2) + + n_heads = q.shape[1] + n_kv = k.shape[1] + enable_gqa = n_heads != n_kv + T = q.shape[2] + T_pad = prefix_seg_info.T_pad + assert T == prefix_seg_info.T, f"Q length {T} != seg.T {prefix_seg_info.T}" + + if block_mask is None: + block_mask = get_block_mask(prefix_seg_info, device, compile_mask = compiled) + + out = _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad) + # back to [1, T, n_heads, D] + return out.transpose(1, 2).contiguous() + + +__all__ = [ + "PrefixSegInfo", + "build_seg_info_multigroup", + "build_seg_info_from_layout", + "get_block_mask", + "clear_block_mask_cache", + "flex_shared_prefix_attention", +]