@bendyline/gezel-inference 0.0.0-stage → 0.1.0

package/dist/providers/mlx/python/template_stability.py

@@ -0,0 +1,133 @@
+"""Make a model's chat template produce a cache-stable prompt prefix.
+
+The problem
+-----------
+Prompt-prefix reuse needs the leading tokens to be byte-identical across
+turns. Qwen 3.5+ templates emit the reasoning-effort preamble as the FIRST
+thing inside the system block, above the tool schemas and above the gezel's
+about/project text::
+
+    <|im_start|>system\\n
+    Reasoning effort is set to xhigh. ...      <- changes per turn
+    # Tools ... <tools> ... </tools>           <- stable per tool roster
+    <gezel about / project docs>               <- stable per gezel+project
+    <|im_end|>
+
+Gezel varies reasoning per turn on purpose — a constrained "write the file
+now" turn suppresses thinking, an ordinary turn does not — so that first
+line changes shape between consecutive turns of the same session. Measured
+on qwen3.8-27b-q4, a thinking-on and a thinking-off render of the SAME
+messages share **3 tokens**. Everything downstream, including a multi-
+thousand-token tool block, is re-prefilled every time the mode flips.
+
+Moving the preamble below the stable content fixes it: the same pair then
+shares 3,029 of 3,071 tokens (99%).
+
+Why a transform and not a curated template
+------------------------------------------
+A hand-maintained copy of each model's jinja goes stale silently the next
+time the model ships a template revision, and it would have to be written
+per model. This applies a narrow, anchored edit to whatever template the
+model actually shipped, then VERIFIES the result behaves — and falls back
+to the original template on any doubt. A prompt that renders wrong is far
+worse than a prompt that caches badly.
+
+Verification checks the property we want rather than the mechanism:
+  * two renders that differ only in reasoning mode share a long prefix, and
+  * the patched render contains the same content as the original (same
+    token multiset), so nothing was dropped or duplicated.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from typing import Any, Callable, Optional, Tuple
+
+# The preamble emission at the top of the tools branch, and the tail where
+# the system content is closed. Anchored on exact source so a template that
+# has been revised upstream simply fails to match and keeps its original
+# behaviour instead of being edited on a guess.
+_EARLY_EMIT = """    {%- if reasoning_instructions %}
+        {{- reasoning_instructions + '\\n\\n' }}
+    {%- endif %}
+    {{- "# Tools"""
+
+_EARLY_REPLACEMENT = """    {{- "# Tools"""
+
+_TAIL = """    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\\n\\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\\n' }}"""
+
+_TAIL_REPLACEMENT = """    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\\n\\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {%- if reasoning_instructions %}
+        {{- '\\n\\n' + reasoning_instructions }}
+    {%- endif %}
+    {{- '<|im_end|>\\n' }}"""
+
+
+def relocate_reasoning_preamble(template_src: str) -> Optional[str]:
+    """Move the reasoning preamble below the tool + system content.
+
+    Returns the rewritten template, or None when the anchors do not match
+    exactly once each — meaning this template is not the shape we know how
+    to edit, and the caller must keep the original.
+    """
+    if not template_src:
+        return None
+    if template_src.count(_EARLY_EMIT) != 1 or template_src.count(_TAIL) != 1:
+        return None
+    patched = template_src.replace(_EARLY_EMIT, _EARLY_REPLACEMENT, 1)
+    patched = patched.replace(_TAIL, _TAIL_REPLACEMENT, 1)
+    return patched if patched != template_src else None
+
+
+def _lcp(a, b) -> int:
+    n = 0
+    for x, y in zip(a, b):
+        if x != y:
+            break
+        n += 1
+    return n
+
+
+def verify_prefix_stability(
+    render: Callable[[str, dict], str],
+    encode: Callable[[str], Any],
+    patched: str,
+    original: str,
+    min_gain: int = 64,
+) -> Tuple[bool, str]:
+    """Prove the patched template is both SAFER and no less complete.
+
+    `render(template, kwargs)` must apply `template` to a fixed probe
+    conversation; `encode` tokenizes. Returns (ok, reason).
+    """
+    on = {"enable_thinking": True, "reasoning_effort": "xhigh"}
+    off = {"enable_thinking": False, "reasoning_effort": "low"}
+    try:
+        p_on, p_off = render(patched, on), render(patched, off)
+        o_on, o_off = render(original, on), render(original, off)
+    except Exception as exc:  # noqa: BLE001 — a template that raises is unusable
+        return False, f"patched template failed to render: {exc}"
+
+    # 1. Nothing lost or duplicated: same tokens, only reordered.
+    if Counter(encode(p_on)) != Counter(encode(o_on)):
+        return False, "patched render changed content, not just order"
+    if Counter(encode(p_off)) != Counter(encode(o_off)):
+        return False, "patched render changed content (thinking-off), not just order"
+
+    # 2. The actual goal: a reasoning-mode flip must stop nuking the prefix.
+    before = _lcp(encode(o_on), encode(o_off))
+    after = _lcp(encode(p_on), encode(p_off))
+    if after < before + min_gain:
+        return False, f"no prefix gain (before={before} after={after})"
+    return True, f"prefix across reasoning modes {before} -> {after} tokens"

package/dist/providers/mlx/python/think_budget.py

@@ -0,0 +1,113 @@
+"""Thinking-budget enforcement for MLX generation.
+
+llama-server enforces `tuning.reasoning.thinkingBudget` via its
+`--reasoning-budget` flag (forced `</think>` closure). The MLX tuning map
+dropped the value (`'reasoning.thinkingBudget': null`), so qwen-family
+models on MLX think UNBOUNDED — wild-caught 2026-08-05: qwen3.6-27b spent
+44,644 chars (~11K tokens) inside one think block with 70 hesitation
+markers ("Wait, ...") against a configured budget of 4,096, and the trial
+still failed. Sampling changes (temp 0.6, mild rp) measurably did NOT
+reduce the rumination; a hard budget is the lever that works on llama.
+
+Mechanism: a logits processor (same hook as the tool grammar) tracks
+whether generation is inside a think block by watching the generated
+token ids. Once the block exceeds the budget, every token except
+`</think>` is masked, forcing closure exactly like llama-server. Each
+think block gets its own budget (interleaved thinking re-opens are
+re-budgeted, not exempted).
+
+Safety posture mirrors SafeToolGrammarProcessor: anything unexpected
+disables enforcement for the rest of the turn rather than killing it.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class ThinkBudgetProcessor:
+    def __init__(
+        self,
+        open_id: int,
+        close_id: int,
+        budget: int,
+        opens_in_think: bool,
+    ) -> None:
+        self.open_id = open_id
+        self.close_id = close_id
+        self.budget = max(1, int(budget))
+        self.in_think = opens_in_think
+        self.count = 0
+        self.disabled = False
+        self.forced = False
+
+    def __call__(self, input_ids: Any, logits: Any) -> Any:
+        if self.disabled:
+            return logits
+        try:
+            last: Optional[int] = None
+            if input_ids is not None and len(input_ids) > 0:
+                last = int(input_ids[-1])
+            if last is not None:
+                if last == self.open_id:
+                    self.in_think = True
+                    self.count = 0
+                elif last == self.close_id:
+                    self.in_think = False
+                elif self.in_think:
+                    self.count += 1
+            if self.in_think and self.count >= self.budget:
+                import mlx.core as mx
+
+                if not self.forced:
+                    self.forced = True
+                    print(
+                        f"[think-budget] budget {self.budget} reached — forcing </think>",
+                        flush=True,
+                    )
+                masked = mx.full(logits.shape, -1e9, dtype=logits.dtype)
+                masked[..., self.close_id] = 0.0
+                return masked
+            return logits
+        except Exception as exc:  # noqa: BLE001 - defensive, mirror grammar posture
+            print(
+                f"[think-budget] processor error; disabling for the rest of this "
+                f"turn: {exc}",
+                flush=True,
+            )
+            self.disabled = True
+            return logits
+
+
+def build_think_budget_processor(
+    tokenizer: Any,
+    budget: Any,
+    opens_in_think: bool,
+) -> Optional[ThinkBudgetProcessor]:
+    """Resolve `<think>`/`</think>` as single special tokens; None when the
+    model has no single-token think tags (gemma's channel format, plain
+    instruct models) or the budget is absent/invalid — enforcement simply
+    doesn't apply there."""
+    try:
+        b = int(budget)
+    except (TypeError, ValueError):
+        return None
+    if b <= 0:
+        return None
+    try:
+        open_ids = tokenizer.encode("<think>", add_special_tokens=False)
+        close_ids = tokenizer.encode("</think>", add_special_tokens=False)
+    except Exception:  # noqa: BLE001
+        return None
+    if len(open_ids) != 1 or len(close_ids) != 1:
+        return None
+    print(
+        f"[think-budget] armed budget={b} opens_in_think={opens_in_think}",
+        flush=True,
+    )
+    return ThinkBudgetProcessor(
+        open_id=int(open_ids[0]),
+        close_id=int(close_ids[0]),
+        budget=b,
+        opens_in_think=opens_in_think,
+    )

package/dist/providers/mlx/python/tool_args_json.py

@@ -0,0 +1,297 @@
+"""Incremental qwen-markup -> JSON argument conversion.
+
+The problem
+-----------
+Qwen's tool body is markup, not JSON::
+
+    <function=write_file><parameter=path>a.html</parameter>
+    <parameter=content>&lt;!doctype html&gt;...</parameter></function>
+
+OpenAI's wire shape needs `function.arguments` to be a JSON *string*, so a
+client that receives the raw markup rejects every call ("Missing required
+fields: `path`, `content`"). mlx_vlm's `qwen3_coder.parse_tool_call` does the
+conversion, but only over a COMPLETE call — which is why streaming raw markup
+and buffering-then-parsing are both dead ends: the first is invalid, the
+second gives no turn boundary and so no early abort.
+
+The approach
+------------
+Convert as the bytes arrive. The saving grace is that only STRING parameters
+are large — `content` is the whole file — and a string can be JSON-escaped
+chunk by chunk. Everything else (ints, floats, bools, objects, arrays) is
+small and needs the whole value for `_convert_param_value`-style coercion, so
+those are buffered per-parameter and emitted when their `</parameter>` lands.
+
+Semantics are copied from mlx_vlm 0.6.6 `tool_parsers/qwen3_coder.py` and must
+stay identical; `test_tool_args_json.py` asserts that by differential-testing
+every case against the reference parser. In particular:
+
+  * exactly ONE leading and ONE trailing "\\n" is stripped from a value
+  * a literal "null" (any case) becomes JSON null
+  * a parameter absent from the schema is treated as a string
+  * unknown/odd types fall through to `ast.literal_eval`, as upstream does
+
+Chunk boundaries are the hazard throughout: `</parameter>`, `</function>` and
+the trailing-newline rule all need lookahead, so the encoder holds back the
+minimum bytes that could still turn out to be part of a marker.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+_STRING_TYPES = {"string", "str", "text", "varchar", "char", "enum"}
+_BOOL_TYPES = {"boolean", "bool", "binary"}
+_OBJ_TYPES = {"object", "array", "arr"}
+
+_PARAM_OPEN = "<parameter="
+_PARAM_CLOSE = "</parameter>"
+_FUNC_CLOSE = "</function>"
+
+
+def arguments_config(func_name: str, tools: Optional[Any]) -> Dict[str, Any]:
+    """Schema `properties` for a tool, mirroring `_get_arguments_config`."""
+    if not tools:
+        return {}
+    for tool in tools:
+        function = tool.get("function") if isinstance(tool, dict) else None
+        if not function:
+            continue
+        if function.get("name") == func_name:
+            params = function.get("parameters")
+            if not params:
+                return {}
+            return params.get("properties", {}) or {}
+    return {}
+
+
+def is_string_param(param_name: str, param_config: Dict[str, Any]) -> bool:
+    """Can this parameter stream, or must its value be buffered whole?
+
+    Unknown parameters count as strings — upstream returns the raw value when
+    the schema has no entry, and raw-string is the streamable case.
+    """
+    param = param_config.get(param_name)
+    if not param:
+        return True
+    declared = param.get("type")
+    if declared is None:
+        return True
+    return str(declared).strip().lower() in _STRING_TYPES
+
+
+def convert_param_value(value: str, param_name: str, param_config: Dict[str, Any]) -> Any:
+    """Coerce a complete value, mirroring `_convert_param_value`."""
+    if value.lower() == "null":
+        return None
+    param = param_config.get(param_name)
+    if not param:
+        return value
+    param_type = str(param["type"]).strip().lower() if "type" in param else "string"
+    if param_type in _STRING_TYPES:
+        return value
+    if param_type.startswith(("int", "uint", "long", "short", "unsigned")):
+        return int(value)
+    if param_type.startswith(("num", "float")):
+        as_float = float(value)
+        as_int = int(as_float)
+        return as_float if (as_float - as_int) != 0 else as_int
+    if param_type in _BOOL_TYPES:
+        return value.lower() == "true"
+    if param_type in _OBJ_TYPES or param_type.startswith(("dict", "list")):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError:
+            return ast.literal_eval(value)
+    return ast.literal_eval(value)
+
+
+def _json_fragment(text: str) -> str:
+    """Escape `text` for the inside of a JSON string literal."""
+    return json.dumps(text)[1:-1]
+
+
+def _longest_partial_suffix(text: str, marker: str) -> int:
+    """Length of the trailing run of `text` that could begin `marker`."""
+    limit = min(len(text), len(marker) - 1)
+    for size in range(limit, 0, -1):
+        if marker.startswith(text[-size:]):
+            return size
+    return 0
+
+
+@dataclass
+class ToolArgsJsonEncoder:
+    """Turn a stream of qwen tool-body text into a stream of JSON argument text.
+
+    Feed the body that follows `<function=NAME>`; `feed()` returns JSON text
+    to append to `function.arguments`. `finish()` closes the object.
+    """
+
+    param_config: Dict[str, Any] = field(default_factory=dict)
+    _buf: str = ""
+    _started: bool = False
+    _closed: bool = False
+    _in_value: bool = False
+    _value_is_string: bool = False
+    _value_name: str = ""
+    _value_buf: str = ""
+    _emitted_value_chars: int = 0
+    _pending_newline: bool = False
+    _first_param: bool = True
+    _leading_stripped: bool = False
+    _opened_quote: bool = False
+
+    def _open(self) -> str:
+        if self._started:
+            return ""
+        self._started = True
+        return "{"
+
+    def feed(self, text: str) -> str:
+        if self._closed:
+            return ""
+        out: List[str] = []
+        self._buf += text
+        while self._buf:
+            if not self._in_value:
+                # Between parameters: wait for the next <parameter= or the
+                # closing </function>.
+                p = self._buf.find(_PARAM_OPEN)
+                f = self._buf.find(_FUNC_CLOSE)
+                if p < 0 and f < 0:
+                    keep = max(
+                        _longest_partial_suffix(self._buf, _PARAM_OPEN),
+                        _longest_partial_suffix(self._buf, _FUNC_CLOSE),
+                    )
+                    self._buf = self._buf[len(self._buf) - keep :] if keep else ""
+                    break
+                if f >= 0 and (p < 0 or f < p):
+                    self._buf = ""
+                    out.append(self._open())
+                    out.append("}")
+                    self._closed = True
+                    break
+                close = self._buf.find(">", p + len(_PARAM_OPEN))
+                if close < 0:
+                    self._buf = self._buf[p:]
+                    break
+                name = self._buf[p + len(_PARAM_OPEN) : close]
+                self._buf = self._buf[close + 1 :]
+                out.append(self._open())
+                if not self._first_param:
+                    out.append(",")
+                self._first_param = False
+                self._value_name = name
+                self._value_is_string = is_string_param(name, self.param_config)
+                self._in_value = True
+                self._value_buf = ""
+                self._emitted_value_chars = 0
+                self._pending_newline = False
+                self._leading_stripped = False
+                self._opened_quote = False
+                # NOTE: the opening quote is deliberately NOT emitted yet.
+                # `_convert_param_value` maps a whole value of "null" to JSON
+                # null BEFORE consulting the schema, so even a string-typed
+                # parameter can turn out to be null — and by then a quote
+                # would already be on the wire. Hold the first few bytes.
+                continue
+
+            end = self._buf.find(_PARAM_CLOSE)
+            chunk = self._buf if end < 0 else self._buf[:end]
+            if end < 0:
+                hold = _longest_partial_suffix(chunk, _PARAM_CLOSE)
+                if hold:
+                    chunk = chunk[: len(chunk) - hold]
+                self._buf = self._buf[len(chunk) :]
+            else:
+                self._buf = self._buf[end + len(_PARAM_CLOSE) :]
+
+            if not self._value_is_string:
+                self._value_buf += chunk
+                if end < 0:
+                    break
+                value = self._value_buf
+                if value.startswith("\n"):
+                    value = value[1:]
+                if value.endswith("\n"):
+                    value = value[:-1]
+                out.append(f"{json.dumps(self._value_name)}:")
+                out.append(json.dumps(convert_param_value(value, self._value_name, self.param_config)))
+                self._in_value = False
+                continue
+
+            # String value. Three lookahead rules interact here:
+            #   * strip exactly one leading newline
+            #   * strip exactly one trailing newline (so the final newline is
+            #     held until we know whether anything follows it)
+            #   * a complete value of "null" is JSON null, not a string, so
+            #     the opening quote waits until the value is longer than
+            #     "null" or the parameter closes
+            if not self._leading_stripped:
+                if chunk.startswith("\n"):
+                    chunk = chunk[1:]
+                self._leading_stripped = True
+            if end >= 0 and chunk.endswith("\n"):
+                chunk = chunk[:-1]
+            if not self._opened_quote:
+                self._value_buf += chunk
+                if end < 0 and len(self._value_buf) <= 4:
+                    break
+                if end >= 0 and self._value_buf.lower() == "null":
+                    out.append(f"{json.dumps(self._value_name)}:null")
+                    self._in_value = False
+                    continue
+                out.append(f"{json.dumps(self._value_name)}:\"")
+                self._opened_quote = True
+                chunk = self._value_buf
+                self._value_buf = ""
+            if self._pending_newline and chunk:
+                out.append(_json_fragment("\n"))
+                self._emitted_value_chars += 1
+                self._pending_newline = False
+            if end < 0 and chunk.endswith("\n"):
+                chunk = chunk[:-1]
+                self._pending_newline = True
+            if chunk:
+                out.append(_json_fragment(chunk))
+                self._emitted_value_chars += len(chunk)
+            if end < 0:
+                break
+            # A held newline at close was the trailing one: drop it.
+            self._pending_newline = False
+            out.append('"')
+            self._in_value = False
+        return "".join(out)
+
+    def finish(self) -> str:
+        """Close whatever is open, so a truncated call still yields JSON."""
+        if self._closed:
+            return ""
+        out: List[str] = [self._open()]
+        if self._in_value:
+            if self._value_is_string:
+                if not self._opened_quote:
+                    out.append(f"{json.dumps(self._value_name)}:\"")
+                    out.append(_json_fragment(self._value_buf))
+                out.append('"')
+            else:
+                value = self._value_buf
+                if value.startswith("\n"):
+                    value = value[1:]
+                if value.endswith("\n"):
+                    value = value[:-1]
+                out.append(f"{json.dumps(self._value_name)}:")
+                try:
+                    out.append(
+                        json.dumps(convert_param_value(value, self._value_name, self.param_config))
+                    )
+                except Exception:
+                    out.append(json.dumps(value))
+            self._in_value = False
+        out.append("}")
+        self._closed = True
+        return "".join(out)