scripts/hierarchical_classify.py

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

import argparse
import json
import os
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

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


def score_nodes(
    nodes: list[dict[str, Any]],
    *,
    base_url: str,
    model: str,
    instruct: str,
    document: str,
    document_type: str,
    threshold: float,
    workers: int,
) -> list[dict[str, Any]]:
    def task(node: dict[str, Any]) -> dict[str, Any]:
        score, _ = score_policy(
            base_url=base_url,
            model=model,
            instruct=instruct,
            query=query_for(node, document_type),
            document=document,
            threshold=threshold,
        )
        return node_result(node, score, document_type=document_type)

    if workers <= 1:
        return [task(node) for node in nodes]
    results: list[dict[str, Any]] = []
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(task, node): node for node in nodes}
        for future in as_completed(futures):
            results.append(future.result())
    return sorted(results, key=lambda item: item["id"])


def relationship_maps(taxonomy: dict[str, Any]) -> tuple[dict[str, list[str]], dict[str, list[str]], dict[str, str]]:
    sub_to_leaves: dict[str, list[str]] = {}
    sc_to_subs: dict[str, list[str]] = {}
    leaf_to_sub: dict[str, str] = {}
    for sc in taxonomy["superclasses"]:
        sc_to_subs[sc["id"]] = [sub["id"] for sub in sc["children"]]
        for sub in sc["children"]:
            sub_to_leaves[sub["id"]] = [leaf["id"] for leaf in sub["children"]]
            for leaf in sub["children"]:
                leaf_to_sub[leaf["id"]] = sub["id"]
    return sub_to_leaves, sc_to_subs, leaf_to_sub


def reconcile_hierarchy(
    evaluated: list[dict[str, Any]], taxonomy: dict[str, Any]
) -> dict[str, Any]:
    """Reconcile raw node decisions into descendant-supported hierarchy decisions.

    Shieldstral scores each policy independently. A broad superclass/subcategory may
    therefore fire lexically even when none of its leaves match. The validated view
    retains a higher-level node only when a matched descendant leaf supports it.
    Raw model decisions are preserved verbatim for auditability.
    """
    by_id = {item["id"]: item for item in evaluated}
    raw_matches = [item for item in evaluated if item.get("matched")]
    matched_leaves = sorted(
        [item for item in raw_matches if item.get("level") == "leaf"],
        key=lambda item: (-float(item["score"]), item["id"]),
    )
    matched_leaf_ids = {item["id"] for item in matched_leaves}
    sub_to_leaves, sc_to_subs, leaf_to_sub = relationship_maps(taxonomy)

    supported_sub_ids = {
        sub_id
        for sub_id, leaf_ids in sub_to_leaves.items()
        if matched_leaf_ids.intersection(leaf_ids)
    }
    supported_sc_ids = {
        sc_id
        for sc_id, sub_ids in sc_to_subs.items()
        if supported_sub_ids.intersection(sub_ids)
    }

    validated_matches: list[dict[str, Any]] = []
    for item in raw_matches:
        level = item.get("level")
        if level == "leaf":
            validated_matches.append(item)
        elif level == "subcategory" and item["id"] in supported_sub_ids:
            validated_matches.append(item)
        elif level == "superclass" and item["id"] in supported_sc_ids:
            validated_matches.append(item)

    validated_ids = {item["id"] for item in validated_matches}
    orphan_matches = sorted(
        [item for item in raw_matches if item["id"] not in validated_ids],
        key=lambda item: (-float(item["score"]), item["id"]),
    )

    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for leaf in matched_leaves:
        grouped[leaf_to_sub[leaf["id"]]].append(leaf)
    same_parent_conflicts = [
        {
            "subcategory_id": sub_id,
            "leaf_ids": [leaf["id"] for leaf in leaves],
            "leaf_names": [leaf["name"] for leaf in leaves],
        }
        for sub_id, leaves in sorted(grouped.items())
        if len(leaves) > 1
    ]

    primary = matched_leaves[0] if matched_leaves else {
        "id": "SAFE",
        "name": "Safe / no descendant-supported leaf policy",
        "level": "safe",
        "score": 1.0,
    }
    if matched_leaves:
        primary_sub = leaf_to_sub[primary["id"]]
        primary_sc = next(
            sc_id for sc_id, sub_ids in sc_to_subs.items() if primary_sub in sub_ids
        )
        validated_path = {
            "superclass": by_id.get(primary_sc),
            "subcategory": by_id.get(primary_sub),
            "leaf": primary,
        }
    else:
        validated_path = None

    return {
        "raw_model_matches": sorted(
            raw_matches, key=lambda item: (-float(item["score"]), item["id"])
        ),
        "hierarchy_validated_matches": sorted(
            validated_matches,
            key=lambda item: ({"superclass": 1, "subcategory": 2, "leaf": 3}[item["level"]], -float(item["score"]), item["id"]),
        ),
        "orphan_policy_matches": orphan_matches,
        "matched_leaf_categories": matched_leaves,
        "primary_leaf": primary,
        "validated_path": validated_path,
        "raw_overall_unsafe": bool(raw_matches),
        "overall_unsafe": bool(matched_leaves),
        "ambiguity": {
            "is_ambiguous": len(matched_leaves) > 1,
            "matched_leaf_count": len(matched_leaves),
            "ambiguous_leaf_matches": matched_leaves if len(matched_leaves) > 1 else [],
            "same_parent_leaf_conflicts": same_parent_conflicts,
        },
        "hierarchy_consistency": {
            "is_consistent": not orphan_matches,
            "orphan_match_count": len(orphan_matches),
            "rule": "A superclass or subcategory is retained only when at least one descendant leaf is matched.",
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Hierarchical 12/26/52 Shieldstral policy classification.")
    parser.add_argument("document")
    parser.add_argument("--document-type", choices=("prompt", "response"), default="prompt")
    parser.add_argument("--mode", choices=("hierarchical", "exhaustive"), default="hierarchical")
    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("--workers", type=int, default=2)
    parser.add_argument("--top", type=int, default=20, help="Maximum raw ranked scores in output; 0 means all")
    args = parser.parse_args()

    taxonomy = load_taxonomy(args.taxonomy)
    model = discover_model(args.base_url, args.model)
    evaluated: list[dict[str, Any]] = []

    if args.mode == "exhaustive":
        all_nodes: list[dict[str, Any]] = []
        for sc in taxonomy["superclasses"]:
            all_nodes.append(sc)
            for sub in sc["children"]:
                all_nodes.append(sub)
                all_nodes.extend(sub["children"])
        evaluated = score_nodes(
            all_nodes,
            base_url=args.base_url,
            model=model,
            instruct=args.instruct,
            document=args.document,
            document_type=args.document_type,
            threshold=args.threshold,
            workers=args.workers,
        )
    else:
        super_results = score_nodes(
            taxonomy["superclasses"],
            base_url=args.base_url,
            model=model,
            instruct=args.instruct,
            document=args.document,
            document_type=args.document_type,
            threshold=args.threshold,
            workers=args.workers,
        )
        evaluated.extend(super_results)
        matched_sc_ids = {item["id"] for item in super_results if item["matched"]}
        sub_nodes = [
            sub
            for sc in taxonomy["superclasses"]
            if sc["id"] in matched_sc_ids
            for sub in sc["children"]
        ]
        sub_results = score_nodes(
            sub_nodes,
            base_url=args.base_url,
            model=model,
            instruct=args.instruct,
            document=args.document,
            document_type=args.document_type,
            threshold=args.threshold,
            workers=args.workers,
        ) if sub_nodes else []
        evaluated.extend(sub_results)
        matched_sub_ids = {item["id"] for item in sub_results if item["matched"]}
        leaf_nodes = [
            leaf
            for sc in taxonomy["superclasses"]
            for sub in sc["children"]
            if sub["id"] in matched_sub_ids
            for leaf in sub["children"]
        ]
        leaf_results = score_nodes(
            leaf_nodes,
            base_url=args.base_url,
            model=model,
            instruct=args.instruct,
            document=args.document,
            document_type=args.document_type,
            threshold=args.threshold,
            workers=args.workers,
        ) if leaf_nodes else []
        evaluated.extend(leaf_results)

    ranked = sorted(evaluated, key=lambda item: item["score"], reverse=True)
    reconciled = reconcile_hierarchy(evaluated, taxonomy)
    result = {
        "classification_mode": f"policy_adaptive_{args.mode}_hierarchy_v2",
        "taxonomy": {
            "name": taxonomy["registry_name"],
            "counts": taxonomy["counts"],
            "query_disclosure": taxonomy["source"]["disclosure"],
        },
        "model": model,
        "instruct": args.instruct,
        "document": args.document,
        "document_type": args.document_type,
        "threshold": args.threshold,
        "primary_class": reconciled["primary_leaf"],
        "evaluated_node_count": len(evaluated),
        "protocol_valid_for_all_evaluated_nodes": all(item["protocol_valid"] for item in evaluated),
        **reconciled,
        "ranked_scores": ranked if args.top == 0 else ranked[: max(args.top, 0)],
    }
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()