#!/usr/bin/env python3
"""Measure how reliably a model and server emit valid tool calls.

Purpose: the course's tool-calling reliability test. Sends a fixed set of prompts -
    fifteen that should produce exactly one named call and five that should produce
    none - to any OpenAI-compatible endpoint with the same tool list every time, and
    reports six rates: call rate, parse rate, right-tool rate, schema validity,
    argument correctness and false-call rate. Nothing here asks a model for an
    opinion; every check is a comparison against the schema in tool-prompts.json.
    Part 25 reuses this script to compare coding models, and Part 27 uses the
    per-case failures to decide what to fine-tune.
Platform: all (pure Python over HTTP; the server may be on any track or another machine)
Minimum memory: 8 GB on the machine running the model; this script needs almost none
Assumes: Python 3.9 or later and no third-party packages. An OpenAI-compatible
    /v1/chat/completions endpoint reachable at --base-url: llama-server from Part 6,
    Ollama or LM Studio from Part 7, vLLM from Part 9, or the Part 9 gateway. The
    server must already be configured for tool calling (Part 9 covers the per-engine
    flags). tool-prompts.json sits beside this file.

Usage: python3 tool-call-reliability.py --base-url http://127.0.0.1:4000/v1 \
           --model local/chat --repeat 5 --labbook labbook.md
       python3 tool-call-reliability.py --base-url http://127.0.0.1:8080/v1 \
           --model local-chat --tool-choice required --temperature 0
       python3 tool-call-reliability.py --base-url http://127.0.0.1:8080/v1 \
           --model local-chat --no-think --json > reliability.json
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

HERE = Path(__file__).resolve().parent
DEFAULT_PROMPTS = HERE / "tool-prompts.json"

# Fragments that mean "the model tried to call a tool and nothing parsed it back".
# Each one is a real format from a model family in this course's reference set, seen
# as visible text: Hermes-style XML from Qwen, harmony channels from gpt-oss, the
# pythonic form, and LM Studio's default format for models with no native template.
UNPARSED_MARKERS = (
    "<tool_call>",
    "</tool_call>",
    "<|channel|>commentary",
    "functions.",
    "[TOOL_REQUEST]",
    "<function=",
    "<tool_use>",
)
LOOSE_JSON_CALL = re.compile(r'\{\s*"(?:name|tool|function)"\s*:\s*"', re.IGNORECASE)


# --------------------------------------------------------------------------------------
# A JSON Schema check small enough to read, covering what tool schemas actually use
# --------------------------------------------------------------------------------------

JSON_TYPES = {
    "string": str,
    "integer": int,
    "number": (int, float),
    "boolean": bool,
    "array": list,
    "object": dict,
}


def schema_errors(arguments: Any, schema: Dict[str, Any]) -> List[str]:
    """Every way `arguments` fails `schema`, as short human-readable strings.

    Deliberately not a full JSON Schema implementation: it covers object type,
    required, additionalProperties, per-property type, enum and numeric bounds,
    which is everything a tool schema in this course uses. A validator you can
    read is worth more here than one you have to trust.
    """
    errors: List[str] = []
    if not isinstance(arguments, dict):
        return ["arguments are not a JSON object"]

    properties = schema.get("properties", {}) or {}
    for name in schema.get("required", []) or []:
        if name not in arguments:
            errors.append("missing required parameter %s" % name)

    if schema.get("additionalProperties") is False:
        for name in arguments:
            if name not in properties:
                errors.append("invented parameter %s" % name)

    for name, value in arguments.items():
        spec = properties.get(name)
        if not isinstance(spec, dict):
            continue
        wanted = spec.get("type")
        expected = JSON_TYPES.get(wanted) if isinstance(wanted, str) else None
        if expected is not None:
            # JSON has no integer type at the wire level, so a whole float is an integer.
            if wanted == "integer" and isinstance(value, float) and value.is_integer():
                value = int(value)
            if isinstance(value, bool) and wanted != "boolean":
                errors.append("%s is a boolean, expected %s" % (name, wanted))
            elif not isinstance(value, expected):
                errors.append("%s is %s, expected %s" % (name, type(value).__name__, wanted))
        if "enum" in spec and value not in spec["enum"]:
            errors.append("%s=%r is not one of %s" % (name, value, spec["enum"]))
        if isinstance(value, (int, float)) and not isinstance(value, bool):
            if "minimum" in spec and value < spec["minimum"]:
                errors.append("%s is below the minimum" % name)
            if "maximum" in spec and value > spec["maximum"]:
                errors.append("%s is above the maximum" % name)
    return errors


def argument_errors(arguments: Dict[str, Any], expected: Dict[str, Any]) -> List[str]:
    """Check the values a case says the call should carry.

    Three comparisons, which is all the test set needs: `equals` for an exact value
    (numbers compared numerically so 24 and 24.0 agree), `contains` for a
    case-insensitive substring, and `one_of` for a small set of acceptable answers.
    """
    errors: List[str] = []
    for name, rule in (expected or {}).items():
        if name not in arguments:
            errors.append("expected %s in the arguments" % name)
            continue
        got = arguments[name]
        if "equals" in rule:
            want = rule["equals"]
            same = (
                abs(float(got) - float(want)) < 1e-9
                if isinstance(want, (int, float)) and isinstance(got, (int, float))
                and not isinstance(got, bool)
                else got == want
            )
            if not same:
                errors.append("%s=%r, expected %r" % (name, got, want))
        if "contains" in rule and rule["contains"].lower() not in str(got).lower():
            errors.append("%s=%r does not contain %r" % (name, got, rule["contains"]))
        if "one_of" in rule and got not in rule["one_of"]:
            errors.append("%s=%r is not one of %s" % (name, got, rule["one_of"]))
    return errors


# --------------------------------------------------------------------------------------
# Talking to the server
# --------------------------------------------------------------------------------------

def post_chat(base_url: str, payload: Dict[str, Any], api_key: Optional[str],
              timeout: int) -> Dict[str, Any]:
    """One /v1/chat/completions request. Raises RuntimeError with a readable message."""
    url = base_url.rstrip("/") + "/chat/completions"
    body = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(url, data=body, method="POST")
    request.add_header("Content-Type", "application/json")
    if api_key:
        request.add_header("Authorization", "Bearer " + api_key)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError("HTTP %s from %s: %s" % (exc.code, url, detail)) from exc
    except urllib.error.URLError as exc:
        raise RuntimeError("could not reach %s: %s" % (url, exc.reason)) from exc


def first_call(message: Dict[str, Any]) -> Tuple[Optional[str], Optional[str], int]:
    """The name and raw argument string of the first tool call, and how many there were."""
    calls = message.get("tool_calls") or []
    if not calls:
        return None, None, 0
    function = calls[0].get("function", {}) or {}
    return function.get("name"), function.get("arguments"), len(calls)


def looks_like_an_unparsed_call(text: str) -> bool:
    """True when the content carries a tool-call format that nothing extracted.

    This is the single most useful diagnostic in the whole test: it separates
    "the model would not call the tool" from "the model called it and the server's
    parser did not recognise the format", which are fixed in completely different
    places.
    """
    if not text:
        return False
    if any(marker in text for marker in UNPARSED_MARKERS):
        return True
    return bool(LOOSE_JSON_CALL.search(text))


# --------------------------------------------------------------------------------------
# The measurement
# --------------------------------------------------------------------------------------

def run_case(case: Dict[str, Any], tools: List[Dict[str, Any]], args: argparse.Namespace,
             schemas: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
    """Send one case `--repeat` times and count what came back."""
    counts = {
        "attempts": 0, "called": 0, "parsed": 0, "right_tool": 0,
        "schema_valid": 0, "args_correct": 0, "unparsed_in_text": 0,
        "multiple_calls": 0, "errors": 0,
    }
    notes: List[str] = []

    payload_base: Dict[str, Any] = {
        "model": args.model,
        "messages": [
            {"role": "system", "content": args.system},
            {"role": "user", "content": case["prompt"]},
        ],
        "tools": tools,
        "temperature": args.temperature,
        "max_tokens": args.max_tokens,
    }
    if args.tool_choice != "auto":
        payload_base["tool_choice"] = args.tool_choice
    if args.no_think:
        # Documented by Qwen for its own models and accepted by vLLM as a per-request
        # template argument. Servers that do not know the key ignore it.
        payload_base["chat_template_kwargs"] = {"enable_thinking": False}

    for _ in range(args.repeat):
        counts["attempts"] += 1
        try:
            response = post_chat(args.base_url, dict(payload_base), args.api_key, args.timeout)
        except RuntimeError as exc:
            counts["errors"] += 1
            notes.append(str(exc)[:160])
            continue

        message = (response.get("choices") or [{}])[0].get("message", {}) or {}
        name, raw_arguments, how_many = first_call(message)
        content = message.get("content") or ""

        if name is None:
            if looks_like_an_unparsed_call(content):
                counts["unparsed_in_text"] += 1
                notes.append("call format in content: " + content.strip()[:120])
            continue

        counts["called"] += 1
        if how_many > 1:
            counts["multiple_calls"] += 1

        try:
            arguments = json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
            if not isinstance(arguments, dict):
                raise ValueError("arguments are not an object")
        except (TypeError, ValueError) as exc:
            notes.append("unparseable arguments (%s): %r" % (exc, str(raw_arguments)[:100]))
            continue
        counts["parsed"] += 1

        if case.get("expect_tool") is None or name != case["expect_tool"]:
            if case.get("expect_tool") is not None:
                notes.append("called %s, expected %s" % (name, case["expect_tool"]))
            continue
        counts["right_tool"] += 1

        problems = schema_errors(arguments, schemas.get(name, {}))
        if problems:
            notes.append("; ".join(problems[:3]))
            continue
        counts["schema_valid"] += 1

        wrong = argument_errors(arguments, case.get("expect_args", {}))
        if wrong:
            notes.append("; ".join(wrong[:3]))
            continue
        counts["args_correct"] += 1

    # Keep the report short: the first three distinct notes say what went wrong.
    seen: List[str] = []
    for note in notes:
        if note not in seen:
            seen.append(note)
    return {"id": case["id"], "expect_tool": case.get("expect_tool"),
            "counts": counts, "notes": seen[:3]}


def summarise(results: List[Dict[str, Any]]) -> Dict[str, Any]:
    """The six rates, computed over the cases each one is meaningful for."""
    expected = [r for r in results if r["expect_tool"] is not None]
    unexpected = [r for r in results if r["expect_tool"] is None]

    def total(rows: List[Dict[str, Any]], key: str) -> int:
        return sum(r["counts"][key] for r in rows)

    def rate(numerator: int, denominator: int) -> Optional[float]:
        return round(numerator / denominator, 4) if denominator else None

    wanted = total(expected, "attempts")
    called = total(expected, "called")
    return {
        "call_rate": rate(called, wanted),
        "parse_rate": rate(total(expected, "parsed"), called),
        "right_tool_rate": rate(total(expected, "right_tool"), wanted),
        "schema_valid_rate": rate(total(expected, "schema_valid"), wanted),
        "args_correct_rate": rate(total(expected, "args_correct"), wanted),
        "false_call_rate": rate(total(unexpected, "called"), total(unexpected, "attempts")),
        "unparsed_in_text": total(results, "unparsed_in_text"),
        "multiple_calls": total(results, "multiple_calls"),
        "request_errors": total(results, "errors"),
        "cases": len(results),
        "attempts": total(results, "attempts"),
    }


def print_report(results: List[Dict[str, Any]], rates: Dict[str, Any], args: argparse.Namespace) -> None:
    print("\n%-12s %-16s %8s %8s %8s %8s" % ("case", "expected tool", "called", "parsed", "valid", "correct"))
    for row in results:
        counts = row["counts"]
        print("%-12s %-16s %8d %8d %8d %8d" % (
            row["id"], row["expect_tool"] or "(none)",
            counts["called"], counts["parsed"], counts["schema_valid"], counts["args_correct"]))
        for note in row["notes"]:
            print("             %s" % note)

    def show(label: str, value: Optional[float]) -> str:
        return "%-22s %s" % (label, "n/a" if value is None else "%.1f%%" % (value * 100))

    print("\n%s, %s, temperature %s, %d attempts per case" % (
        args.model, args.base_url, args.temperature, args.repeat))
    print(show("call rate", rates["call_rate"]))
    print(show("parse rate", rates["parse_rate"]))
    print(show("right tool", rates["right_tool_rate"]))
    print(show("schema valid", rates["schema_valid_rate"]))
    print(show("arguments correct", rates["args_correct_rate"]))
    print(show("false calls", rates["false_call_rate"]))
    if rates["unparsed_in_text"]:
        print("\n%d response(s) contained a tool-call format as visible text. That is a parser or "
              "chat-template mismatch on the server, not a limitation of the model: check the "
              "engine's tool-call parser against the model family before concluding anything."
              % rates["unparsed_in_text"])
    if rates["request_errors"]:
        print("%d request(s) failed outright; the rates above are computed over the rest."
              % rates["request_errors"])


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", default="http://127.0.0.1:4000/v1",
                        help="OpenAI-compatible base URL, ending in /v1")
    parser.add_argument("--model", required=True, help="the served model name or alias")
    parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"),
                        help="bearer token, if the endpoint needs one; defaults to $OPENAI_API_KEY")
    parser.add_argument("--prompts", default=str(DEFAULT_PROMPTS), help="the test set to run")
    parser.add_argument("--repeat", type=int, default=5, help="attempts per case")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--tool-choice", default="auto", choices=["auto", "required", "none"])
    parser.add_argument("--no-think", action="store_true",
                        help="ask the server to disable thinking mode for this run")
    parser.add_argument("--system", default=(
        "You are a careful assistant with tools. Call a tool when one of them can answer the "
        "request. Answer directly when none of them applies."))
    parser.add_argument("--timeout", type=int, default=180)
    parser.add_argument("--json", action="store_true", help="print the whole report as JSON")
    parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file")
    parser.add_argument("--notes", default=None, help="free text recorded with the run")
    args = parser.parse_args()

    if args.repeat > 1 and args.temperature == 0:
        print("note: at temperature 0 every attempt is the same sample, so --repeat measures "
              "nothing but server determinism.", file=sys.stderr)

    suite = json.loads(Path(args.prompts).read_text(encoding="utf-8"))
    tools = suite["tools"]
    schemas = {t["function"]["name"]: t["function"].get("parameters", {}) for t in tools}

    started = time.time()
    results = []
    for case in suite["cases"]:
        results.append(run_case(case, tools, args, schemas))
        print(".", end="", flush=True)
    elapsed = time.time() - started
    rates = summarise(results)

    if args.json:
        print(json.dumps({"rates": rates, "results": results}, indent=2))
    else:
        print_report(results, rates, args)
        print("\n%d requests in %.0f s" % (rates["attempts"], elapsed))

    if args.labbook:
        record = {
            "lab": "part-24/tool-call-reliability",
            "model": args.model,
            "base_url": args.base_url,
            "prompts": os.path.basename(args.prompts),
            "prompts_version": suite.get("version"),
            "repeat": args.repeat,
            "temperature": args.temperature,
            "tool_choice": args.tool_choice,
            "thinking_disabled": bool(args.no_think),
            "rates": rates,
            "seconds": round(elapsed, 1),
            "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "notes": args.notes,
        }
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\n")
        print("recorded in %s" % args.labbook)


if __name__ == "__main__":
    main()