recipe/pipeline/nex_tools_diag.py

#!/usr/bin/env python3
"""Diagnose the Nex tool-call gate (n-tools-q106: 6/14 - every thinking-on check failed, and nested-object failed
with thinking off on an HTTP 500 "does not match the expected peg-native format"). nex_bench.py keeps only the last
800 characters of the harness output, so the per-check reasons were lost.

Same binary, file, server flags and tool schema as the gate, but every raw reply is kept (content, reasoning_content,
tool_calls, HTTP error body). Variants, one server session:
  gate_on       reasoning_effort=high, gate sampling (temp 0.7, top_p 0.95, top_k 40)   - the failing half
  gate_off      reasoning_effort=none, gate sampling                                    - the passing half
  on_greedy     reasoning_effort=high, temp 0 / top_k 1                                 - is it sampling?
  adaptive      no reasoning_effort (the template's adaptive mode), gate sampling
  nested_off_x3 the nested-object request three more times with thinking off           - is the 500 repeatable?
Diagnostic only (not a card measurement) -> results/nex_tools_diag.json."""
import json, os, sys, time, urllib.error, urllib.request
from types import SimpleNamespace

os.environ["AGNES_BIN"] = "/opt/llama-rocm/rocmfpx-724/build-hipvk/bin"   # the gate ran on the unpatched server
sys.path.insert(0, "/mnt/models/nex-n2.5-mini")
import nex_harness as H  # noqa: E402

W = H.W
PORT = 18650
OUT = f"{W}/results/nex_tools_diag.json"
a = SimpleNamespace(model=f"{W}/out/Nex-N2.5-mini-Q4_0_ROCMFP4_STRIX_LEAN.gguf", dev="ROCm0", ctx=65536,
                    draft=None, mtp_infile=False, nmax=4, pmin=0.0, strict=False,
                    serverlog=f"{W}/logs/diag_tools_server.log")
PROMPTS = {
    "multi-arg": "What's the weather in Paris in celsius?",
    "nested-object": "Book 'Design review' on 2026-10-02 at 14:00 with ana@x.io and bo@x.io.",
    "enum": "Weather in Denver, and give it to me in fahrenheit.",
    "correct-decline": "What is 17 times 23? Answer directly.",
    "multi-turn": "What's the weather in Tokyo in celsius?",
    "streaming": "What's the weather in Rome in celsius?",
    "parallel": "Get the weather in Oslo AND in Lima, both in celsius. Call the tool for each city.",
}
GATE = dict(temperature=0.7, top_p=0.95, top_k=40)
GREEDY = dict(temperature=0, top_k=1)


def call(msgs, kwargs, sampling, stream=False):
    body = {"messages": msgs, "tools": H.TOOLS, "tool_choice": "auto", "max_tokens": 8192, **sampling}
    if kwargs is not None:
        body["chat_template_kwargs"] = kwargs
    t0 = time.time()
    try:
        if not stream:
            r = H.post(PORT, "/v1/chat/completions", body)
            ch = r["choices"][0]
            return {"message": ch["message"], "finish": ch.get("finish_reason"), "s": round(time.time() - t0, 2)}
        body["stream"] = True
        req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", data=json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        calls, content, reasoning, finish = {}, "", "", None
        for line in urllib.request.urlopen(req, timeout=1800):
            line = line.decode().strip()
            if not line.startswith("data:") or line.endswith("[DONE]"):
                continue
            ch = json.loads(line[5:])["choices"][0]
            d = ch.get("delta") or {}
            finish = ch.get("finish_reason") or finish
            content += d.get("content") or ""
            reasoning += d.get("reasoning_content") or ""
            for tc in d.get("tool_calls") or []:
                c = calls.setdefault(tc["index"], {"name": "", "arguments": ""})
                c["name"] += (tc.get("function") or {}).get("name") or ""
                c["arguments"] += (tc.get("function") or {}).get("arguments") or ""
        msg = {"content": content, "reasoning_content": reasoning,
               "tool_calls": [{"function": v} for _, v in sorted(calls.items())]}
        return {"message": msg, "finish": finish, "s": round(time.time() - t0, 2)}
    except urllib.error.HTTPError as e:
        return {"http_error": e.code, "body": e.read().decode(errors="replace")[:600], "s": round(time.time() - t0, 2)}
    except Exception as e:  # noqa: BLE001 - a diagnostic records every failure mode
        return {"exception": repr(e)[:300], "s": round(time.time() - t0, 2)}


def summarize(res):
    if "message" not in res:
        return res
    m = res["message"]
    c = m.get("content") or ""
    return {"finish": res.get("finish"), "s": res["s"],
            "tool_calls": [(t.get("function") or {}).get("name") for t in (m.get("tool_calls") or [])],
            "args": [(t.get("function") or {}).get("arguments") for t in (m.get("tool_calls") or [])],
            "leaks_in_content": [x for x in H.LEAK if x in c],
            "content": c[:600], "reasoning_len": len(m.get("reasoning_content") or ""),
            "reasoning_head": (m.get("reasoning_content") or "")[:300]}


def run_variant(name, kwargs, sampling, stream_ok=True):
    out = {}
    for k, p in PROMPTS.items():
        msgs = [{"role": "user", "content": p}]
        r = call(msgs, kwargs, sampling, stream=(k == "streaming" and stream_ok))
        rec = {"first": summarize(r)}
        if k == "multi-turn" and "message" in r and r["message"].get("tool_calls"):
            tc = r["message"]["tool_calls"][0]
            msgs += [{"role": "assistant", "content": r["message"].get("content") or "",
                      "tool_calls": [{"id": "call_1", "type": "function", "function": tc["function"]}]},
                     {"role": "tool", "tool_call_id": "call_1", "content": json.dumps({"temp_c": 21, "sky": "clear"})}]
            rec["second"] = summarize(call(msgs, kwargs, sampling))
        out[k] = rec
        print(name, k, json.dumps(rec)[:400], flush=True)
    return out


s = H.Server(a, PORT)
report = {"model": os.path.basename(a.model), "bin": H.BIN, "cmd": " ".join(s.cmd), "variants": {}}
try:
    report["variants"]["gate_on"] = run_variant("gate_on", H.THINK_ON, GATE)
    report["variants"]["gate_off"] = run_variant("gate_off", H.THINK_OFF, GATE)
    report["variants"]["on_greedy"] = run_variant("on_greedy", H.THINK_ON, GREEDY, stream_ok=False)
    report["variants"]["adaptive"] = run_variant("adaptive", None, GATE, stream_ok=False)
    report["variants"]["nested_off_x3"] = [
        summarize(call([{"role": "user", "content": PROMPTS["nested-object"]}], H.THINK_OFF, GATE)) for _ in range(3)]
    print("nested_off_x3", json.dumps(report["variants"]["nested_off_x3"])[:600], flush=True)
finally:
    s.stop()
json.dump(report, open(OUT, "w"), indent=1)
print("NEX_TOOLS_DIAG_DONE", OUT)