#!/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()