scripts/model_probe.py

import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import app  # noqa: E402


def _assistant_text(result):
    history = result[0] or []
    return history[-1]["content"] if history else ""


def _scenario(name, message, history, active_schema, state):
    result = app.generate_response(
        message,
        history,
        active_schema,
        app.FINE_TUNED_MODEL_KEY,
        state,
    )
    return {
        "name": name,
        "message": message,
        "assistant": _assistant_text(result),
        "sql": result[4],
        "status": result[6],
        "active_schema": result[2],
        "state": result[7],
        "history": result[0],
    }


def _contains_any(text, needles):
    text = (text or "").lower()
    return any(needle.lower() in text for needle in needles)


def _grade(records):
    checks = []
    by_name = {record["name"]: record for record in records}

    checks.append({
        "name": "smalltalk_is_static_fallback",
        "pass": app.FALLBACK_RESPONSE in by_name["greeting"]["assistant"] and not by_name["greeting"]["sql"],
        "detail": "Greeting should use the static fallback and no SQL.",
    })
    checks.append({
        "name": "create_table_is_deterministic",
        "pass": "CREATE TABLE ZOOLOGICO" in (by_name["create_schema"]["sql"] or "").upper(),
        "detail": "Explicit CREATE TABLE request should not need model generation.",
    })
    checks.append({
        "name": "edit_updates_schema",
        "pass": _contains_any(by_name["edit_schema"]["sql"], ["numero_animais", "num_animais"]),
        "detail": "Edit should replace capacidade with an animal-count column.",
    })
    checks.append({
        "name": "query_generates_select",
        "pass": "SELECT" in (by_name["query_schema"]["sql"] or "").upper(),
        "detail": "Natural query should generate SELECT SQL.",
    })
    checks.append({
        "name": "smalltalk_with_schema_stays_chat",
        "pass": app.FALLBACK_RESPONSE in by_name["smalltalk_with_schema"]["assistant"] and not by_name["smalltalk_with_schema"]["sql"],
        "detail": "Smalltalk with active schema should still use static fallback.",
    })
    return checks


def main():
    app.load_model(app.FINE_TUNED_MODEL_ID)

    history = []
    active_schema = ""
    state = app.chat_core.default_state()
    records = []

    for name, message in [
        ("greeting", "oi"),
        ("create_schema", "crie tabela zoologico com id nome cidade capacidade"),
        ("edit_schema", "troca capacidade por numero_animais"),
        ("query_schema", "liste zoologicos de Sao Paulo"),
        ("smalltalk_with_schema", "como voce esta hoje?"),
    ]:
        record = _scenario(name, message, history, active_schema, state)
        records.append({key: value for key, value in record.items() if key != "history"})
        history = record["history"]
        active_schema = record["active_schema"]
        state = record["state"]

    checks = _grade(records)
    report = {
        "model": app.FINE_TUNED_MODEL_ID,
        "passed": all(check["pass"] for check in checks),
        "checks": checks,
        "records": records,
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0 if report["passed"] else 1


if __name__ == "__main__":
    raise SystemExit(main())