scripts/classify_node.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os

from shieldstral_policy import (
    discover_model,
    flatten_taxonomy,
    load_taxonomy,
    node_result,
    query_for,
    score_policy,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Score one Shieldstral hierarchy policy node.")
    parser.add_argument("--node", required=True, help="Official SC/CAT ID or AXONVERTEX SUB ID")
    parser.add_argument("--document", required=True)
    parser.add_argument("--document-type", choices=("prompt", "response"), default="prompt")
    parser.add_argument("--instruct", default="Evaluate whether the document matches the query criteria")
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--base-url", default=os.getenv("BASE_URL", "http://127.0.0.1:18190/v1"))
    parser.add_argument("--model", default=None)
    parser.add_argument("--taxonomy", default=None)
    parser.add_argument("--raw", action="store_true")
    args = parser.parse_args()

    taxonomy = load_taxonomy(args.taxonomy)
    nodes = flatten_taxonomy(taxonomy)
    node_id = args.node.upper()
    if node_id not in nodes:
        raise SystemExit(f"Unknown node ID: {node_id}")
    node = nodes[node_id]
    model = discover_model(args.base_url, args.model)
    score, raw = score_policy(
        base_url=args.base_url,
        model=model,
        instruct=args.instruct,
        query=query_for(node, args.document_type),
        document=args.document,
        threshold=args.threshold,
    )
    result = {
        "classification_mode": "single_policy_node",
        "model": model,
        "document_type": args.document_type,
        "node": node_result(node, score, document_type=args.document_type),
    }
    if args.raw:
        result["raw_response"] = raw
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()