scripts/check_bridge_sync.py

#!/usr/bin/env python3
"""
Janus-35B — verify Modelfile and HF Ollama bridge files stay in sync.

The repo ships two parallel Ollama configurations:

  - ``Modelfile`` is consumed by the local-build path
    (``ollama create janus -f Modelfile``). It contains
    ``TEMPLATE`` / ``SYSTEM`` / ``PARAMETER`` directives.
  - ``template`` / ``system`` / ``params`` at the repo root are consumed by HF's
    Ollama bridge when users ``ollama run hf.co/FoolDev/Janus-35B-HERETIC`` directly. HF
    does NOT read the Modelfile (per https://huggingface.co/docs/hub/en/ollama).

If the two configurations drift apart, ``hf.co/...`` users and local-build
users get different behaviour — exactly the bug fixed in commit 70ccef1
("Add HF Ollama bridge files (template/system/params)"). This script is
the regression guard: it parses the Modelfile, loads the three bridge
files, and fails on any mismatch.

Usage:
    python3 scripts/check_bridge_sync.py
    # exit 0 if in sync, 1 (with diff details) if not.

Run this manually before pushing a Modelfile / bridge-file edit. The 27B
sibling repo wires an equivalent script into scripts/check.sh and a
pre-commit hook; this repo intentionally stays leaner and runs it
on demand.
"""
from __future__ import annotations

import ast
import json
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

# Ollama Modelfile reference: https://docs.ollama.com/modelfile
TEMPLATE_RE = re.compile(r'^TEMPLATE\s+"""(.*?)"""', re.DOTALL | re.MULTILINE)
SYSTEM_RE = re.compile(r'^SYSTEM\s+"""(.*?)"""', re.DOTALL | re.MULTILINE)
PARAMETER_RE = re.compile(r'^PARAMETER\s+(\S+)\s+(.*?)\s*$', re.MULTILINE)


def parse_modelfile(text: str) -> tuple[str, str, dict[str, object]]:
    """Extract TEMPLATE, SYSTEM, and PARAMETER blocks from a Modelfile."""
    tpl_match = TEMPLATE_RE.search(text)
    if not tpl_match:
        die("Modelfile has no TEMPLATE block")
    template = tpl_match.group(1)

    sys_match = SYSTEM_RE.search(text)
    if not sys_match:
        die("Modelfile has no SYSTEM block")
    system = sys_match.group(1)

    params: dict[str, object] = {}
    stops: list[str] = []
    for key, raw in PARAMETER_RE.findall(text):
        # Strip outer quotes if present.
        value: object = raw.strip()
        if isinstance(value, str) and len(value) >= 2 and value[0] == value[-1] == '"':
            value = value[1:-1]
        # Stop tokens accumulate; everything else is scalar.
        if key == "stop":
            stops.append(value)  # type: ignore[arg-type]
            continue
        # Cast anything numeric so it compares equal to the JSON-parsed
        # bridge value (int() then float() mirrors json.loads). A key
        # whitelist would leave any new numeric param a str and trip a
        # false "params drift" against its JSON number.
        try:
            value = int(value)  # type: ignore[arg-type]
        except (TypeError, ValueError):
            try:
                value = float(value)  # type: ignore[arg-type]
            except (TypeError, ValueError):
                pass
        params[key] = value

    if stops:
        params["stop"] = stops

    return template, system, params


def die(msg: str) -> None:
    print(f"[FAIL] {msg}", file=sys.stderr)
    sys.exit(1)


def diff_strings(label: str, expected: str, actual: str) -> bool:
    if expected == actual:
        return True
    print(f"[FAIL] {label} drift detected", file=sys.stderr)
    print(f"  Modelfile len={len(expected)}  bridge file len={len(actual)}", file=sys.stderr)
    # Show the first diverging line for quick orientation.
    e_lines = expected.splitlines()
    a_lines = actual.splitlines()
    for i, (e, a) in enumerate(zip(e_lines, a_lines)):
        if e != a:
            print(f"  first diff at line {i + 1}:", file=sys.stderr)
            print(f"    modelfile : {e!r}", file=sys.stderr)
            print(f"    bridge    : {a!r}", file=sys.stderr)
            return False
    if len(e_lines) != len(a_lines):
        print(f"  line count differs: modelfile={len(e_lines)} bridge={len(a_lines)}",
              file=sys.stderr)
    return False


def check_examples_system_sync(canonical_system: str) -> bool:
    """Every non-vision examples/*.py that embeds a ``*_SYSTEM`` prompt
    constant must match the root ``system`` file verbatim (outer whitespace
    ignored).

    Guards against a text example client silently drifting from the shipped
    system prompt. The vision client is skipped below — it deliberately uses
    its own image-specific prompt.
    """
    examples = ROOT / "examples"
    if not examples.is_dir():
        return True
    want = canonical_system.strip()
    ok = True
    for path in sorted(examples.glob("*.py")):
        # The vision client uses a deliberately image-specific system prompt,
        # not the shared general-assistant one — exclude it from this check.
        if "vision" in path.name:
            continue
        tree = ast.parse(path.read_text())
        for node in ast.walk(tree):
            if not isinstance(node, ast.Assign):
                continue
            names = [t.id for t in node.targets if isinstance(t, ast.Name)]
            if not any(n.endswith("_SYSTEM") for n in names):
                continue
            try:
                value = ast.literal_eval(node.value)
            except (ValueError, SyntaxError):
                continue
            if not isinstance(value, str) or value.strip() == want:
                continue
            print(f"[FAIL] examples/{path.name}: {'/'.join(names)} drifted "
                  f"from root `system`", file=sys.stderr)
            for i, (w, a) in enumerate(zip(want.splitlines(),
                                           value.strip().splitlines())):
                if w != a:
                    print(f"  first diff at line {i + 1}:", file=sys.stderr)
                    print(f"    system  : {w!r}", file=sys.stderr)
                    print(f"    example : {a!r}", file=sys.stderr)
                    break
            ok = False
    return ok


def main() -> int:
    modelfile = (ROOT / "Modelfile").read_text()
    bridge_template = (ROOT / "template").read_text()
    bridge_system = (ROOT / "system").read_text()
    bridge_params = json.loads((ROOT / "params").read_text())

    mf_template, mf_system, mf_params = parse_modelfile(modelfile)

    ok = True

    # 1. TEMPLATE: byte-for-byte.
    ok &= diff_strings("TEMPLATE", mf_template, bridge_template)

    # 2. SYSTEM: trim trailing whitespace on both ends. The bridge file
    #    typically has a trailing newline; the Modelfile block doesn't.
    ok &= diff_strings("SYSTEM", mf_system.strip(), bridge_system.strip())

    # 3. PARAMETER vs params JSON: compare normalized dicts.
    if mf_params != bridge_params:
        print("[FAIL] params drift detected", file=sys.stderr)
        for k in sorted(set(mf_params) | set(bridge_params)):
            mv = mf_params.get(k, "<missing>")
            bv = bridge_params.get(k, "<missing>")
            if mv != bv:
                print(f"  {k}: modelfile={mv!r}  bridge={bv!r}", file=sys.stderr)
        ok = False

    if not ok:
        print("\n[!] Modelfile and bridge files are out of sync.", file=sys.stderr)
        print("    Edit them together: any change to TEMPLATE / SYSTEM /",
              file=sys.stderr)
        print("    PARAMETER must be reflected in template / system / params.",
              file=sys.stderr)

    # 4. examples/*.py *_SYSTEM constants must match the root `system` file.
    examples_ok = check_examples_system_sync(bridge_system)
    if not examples_ok:
        print("\n[!] An examples/ client's system prompt drifted from `system`.",
              file=sys.stderr)
        print("    Update the *_SYSTEM constant to match the root `system` "
              "verbatim.", file=sys.stderr)

    if not (ok and examples_ok):
        return 1

    print("[ ok ] Modelfile <-> bridge files + examples in sync")
    return 0


if __name__ == "__main__":
    sys.exit(main())