Asago Policy Mapper

Unstructured policy-to-risk mapping via AI Risk Atlas Nexus.

Corporate AI policies exist as unstructured documents — PDFs, Word files, HTML pages — written in natural language. Safety tools, evaluation frameworks, and governance processes all deal with risks, but there is no automated way to bridge from raw policy text to those risks.

This semantic gap — from "the model must not provide medical advice" to atlas-hallucination, nist-ms-2.5, owasp-llm09-2025 — is the critical transformation that enables all downstream automation.

This software takes raw, unstructured policy documents and produces a risk landscape: a set of AI Risk Atlas Nexus risk identifiers with enrichments.

• Risk identification — Identifies Nexus risk IDs (e.g., atlas-hallucination, air-2024-0042, nist-ms-2.5) directly from policy text
• Cross-taxonomy mapping — A static SSSOM mapping (src/asago_policy_mapper/data/risk_to_category.sssom.tsv) maps extracted risks to category-level taxonomies (NIST AI RMF, OWASP Top 10 LLM, OWASP ASI, AILuminate)
• Evidence grounding — Each identified risk is grounded to specific passages in the source document, providing traceability from risk to policy text
• Confidence scoring — Each risk mapping includes a confidence score, enabling human review of uncertain mappings

# Input: "acme-ai-policy.pdf"
#   contains: "The AI system must not generate content that could
#              be construed as medical advice..."

# Output:
risk_extraction:
  risks:
    - nexus_id: atlas-hallucination
      confidence: 0.92
      evidence:
        - exact: "The AI system must not generate content that could be construed as medical advice"
          document: "acme-ai-policy.pdf"
          page: 12
      cross_mappings:
        - nist-ms-2.5 (Confabulation)
        - owasp-llm09-2025 (Misinformation)
        - air-2024-0156 (Health misinformation)

The service parses and chunks input documents, then uses hybrid retrieval (keyword and semantic search) against the Nexus risk catalogue to identify candidate risks directly from the policy text. By default, an LLM generates search queries from chunk groups in risk-taxonomy vocabulary (disable with --no-query-gen to use raw chunk text). An LLM then grounds accepted candidates with evidence spans.

flowchart LR
    docs@{ shape: doc, label: "Documents" } --> parse[Parse]
    parse --> chunk[Chunk]
    chunk --> filter[Filter\nagentic]
    filter --> index[Index]
    chunk --> qgen@{ shape: st-rect, label: "Query gen" }
    qgen --> retrieve[Retrieve]
    index --> retrieve
    chunk -.->|--no-query-gen| retrieve
    retrieve --> judge@{ shape: st-rect, label: "Judge" }
    retrieve -.->|--no-judge| ground
    judge --> ground@{ shape: st-rect, label: "Ground" }
    ground --> vg@{ shape: st-rect, label: "Variant\nground" }
    retrieve -.->|--no-grounding| merge
    vg --> merge[Merge]
    merge --> expand@{ shape: st-rect, label: "Expand" }
    expand --> causal@{ shape: st-rect, label: "Causal\nsynthesis" }
    expand -.->|--no-causal-synthesis| miti
    causal -.->|CLI post‑processing| miti[Mitigations]

1. Parse — Docling converts PDF/DOCX/HTML to markdown
2. Chunk — Split into ~512-token chunks with page/section metadata
3. Filter agentic risks — If the document does not contain agent-related terminology, agentic risks are removed from the catalogue before indexing
4. Index — Build BM25 + bi-encoder + cross-encoder index over Nexus risks. Variant risks (IDs containing ---) are collapsed into synthetic parent entries for indexing.
5. Query gen (optional, on by default) — LLM generates 1-3 search queries per section group in risk-taxonomy vocabulary (parallel; disable with --no-query-gen)
6. Retrieve — Hybrid search using generated queries (or, with --no-query-gen, per-chunk BM25 + semantic search with RRF fusion and cross-encoder reranking)
7. Judge — LLM judges borderline candidates (parallel; only applies to fallback chunks in query-gen mode, since query-gen chunks bypass borderline classification)
8. Ground — LLM extracts evidence passages and confidence (parallel, multi-pass)
9. Variant ground — For collapsed parent risks that survived grounding, a specialized LLM call selects only the specifically evidenced variant sub-types
10. Merge — Deduplicate across chunks, keep top-3 evidence spans
11. Expand — Sibling expansion: found risks are expanded to parent siblings + cross-taxonomy mappings, then grounded against relevant document chunks (parallel)
12. Causal synthesis — LLM synthesizes threat-source → threat → vulnerability → consequence → impact chains per matched risk (parallel; disable with --no-causal-synthesis)
13. Mitigations — CLI post-processing: enriches results with mitigation actions and risk cross-maps from Nexus (not part of the extraction pipeline)

The pipeline extracts risk-level risks (IBM Risk Atlas, Credo UCF, AIR 2024, MIT AI Risk Repository — ~486 specific risks). Evaluation runs at two tiers:

• Tier 1 (risk-level): precision/recall/F1 on exact risk ID matches against ground truth
• Tier 2 (category-level): risk IDs are mapped to category-level taxonomies (NIST AI RMF, OWASP Top 10 LLM, OWASP ASI) via a static SSSOM cross-taxonomy mapping (src/asago_policy_mapper/data/risk_to_category.sssom.tsv), then precision/recall/F1 is computed per category taxonomy

Category-level eval answers "did we find the right risk themes?" — more forgiving than risk-level since finding any bias-related risk satisfies the NIST harmful-bias-or-homogenization category.

src/asago_policy_mapper/data/risk_to_category.sssom.tsv is a static SSSOM file mapping 486 risk-level risks to 4 category-level taxonomies (NIST AI RMF 12 risks, OWASP LLM 10 risks, AILuminate 12 risks, OWASP ASI 10 risks). Built from Nexus mapping files + manually reviewed gap-fill for IBM agentic risks, Credo, MIT, and AIR 2024 (314 risks via group-level inheritance). Contains 802 entries; only strong predicates (exact/close/broadMatch) are used at eval time — relatedMatch is excluded.

src/asago_policy_mapper/data/atlas_risk_to_actions.yaml maps 80 Atlas risk IDs to ~552 recommended mitigation actions across 3 frameworks:

• OWASP LLM Top 10 v2.0 (114 action-risk links) — src/asago_policy_mapper/data/owasp_llm_2.0_actions_data.yaml
• NIST AI RMF 600-1 (338 action-risk links) — src/asago_policy_mapper/data/nist_ai_rmf_actions_to_atlas_data.yaml
• AIUC-1 (100 action-risk links) — src/asago_policy_mapper/data/aiuc1_actions_to_atlas_data.yaml

All mappings are direct hasRelatedRisk: atlas-* — no transitive cross-framework hops. Each action is categorized as technical, operational, or governance via rules in src/asago_policy_mapper/data/mitigation_categories.yaml. Regenerate after data file changes: python scripts/build_mitigation_index.py

LLM prompts are Jinja2 templates in src/asago_policy_mapper/templates/prompts/ (rendered by prompts.py::render_prompt()). Each prompt requires a *_user.j2 template and may optionally define a *_system.j2 template; current prompt names with user templates are judge_risk, judge_risk_gepa_demos, generate_queries, ground_evidence, ground_variants, ground_group, and causal_synthesis.

All LLM calls go through llm.py, which wraps the OpenAI client with instructor for structured Pydantic outputs. TokenTracker accumulates token usage across pipeline stages; LLMConfig holds connection details.

• Automatic retry on validation errors (appends error hint), context overflow detection (reduces the configured output-token parameter), and prompt truncation on incomplete output
• LLMConfig.max_tokens (default 8192) is always sent on chat completions. Without this, vLLM treats omitted max_tokens as “fill the remaining context”, so grounding calls can generate for tens of minutes. Override the budget with --max-tokens, POLICY_MAPPER_MAX_OUTPUT_TOKENS, or --max-context.
• The output-token parameter defaults to max_tokens for vLLM compatibility. For endpoints that require OpenAI's newer parameter, set --output-token-parameter max_completion_tokens or POLICY_MAPPER_OUTPUT_TOKEN_PARAMETER=max_completion_tokens. Only the selected parameter is sent.
• An optional processing tier can be passed unchanged to the LLM endpoint with --service-tier or POLICY_MAPPER_SERVICE_TIER. For OpenAI, current examples include default (Standard), flex, and fast/priority (Fast mode); other API-compatible endpoints may support different values. If omitted, no parameter is sent and the endpoint/project default is used. This applies to LLM chat-completion calls; Batch API jobs are not currently supported by the mapper.
• Grounding splits candidate lists into --grounding-batch-size risks per call (default 15) so structured JSON fits in the output budget.
• Sampling parameters (temperature, top_p, top_k) are injected by the tracking wrapper from LLMConfig defaults — call sites don't set them directly. top_k is passed via extra_body for vLLM compatibility.
• All LLM calls default to temperature=0.0; override with --temperature, --top-p, --top-k CLI flags (e.g. --temperature 1.0 --top-p 0.95 --top-k 64 for Gemma 4)

The pipeline uses three models: a bi-encoder for initial semantic retrieval, a cross-encoder for reranking, and an LLM for judging and grounding. The defaults run locally without GPU, but quality improves significantly with better models served remotely via vLLM.

F1 scores are from IR-only evaluation (no LLM judge/grounding) on 27 policies. With LLM stages enabled, the local default achieves F1=0.719 end-to-end; with sibling expansion (--expand-siblings), Qwen3+GTE achieves F1=0.753.

Qwen3-Embedding-4B (recommended) — instruction-aware, 2560-dim, 8K context. Best first-stage retrieval: higher precision than other bi-encoders at comparable recall. Requires a query instruction and remote serving via vLLM.

google/EmbeddingGemma-300M — 300M params, good quality without instructions. Slightly better recall than mpnet (0.895 vs 0.871).

all-mpnet-base-v2 (default) — 110M params, runs locally on CPU. Good baseline but instruction-unaware.

Alibaba-NLP/gte-reranker-modernbert-base (recommended) — AUC=0.759 on pipeline-mined negatives. Genuinely discriminates relevant from irrelevant candidates. Outputs calibrated scores (no sigmoid needed). Serve via vLLM's /v1/score endpoint.

cross-encoder/ms-marco-MiniLM-L-12-v2 (default) — AUC=0.498 on pipeline-mined negatives (essentially random). Functions as a volume reduction filter rather than a semantic discriminator. Runs locally. Works well enough end-to-end because the LLM grounding stage provides the actual precision filtering.

ColBERT late-interaction models are supported via --colbert-model (replaces bi-encoder + cross-encoder with a single model using MaxSim scoring). ColBERT models are local-only — vLLM returns pooled embeddings, not token-level.

Pass a model name to run locally (downloaded on first use), or a URL to use a remote model served via vLLM.

Authentication: Remote endpoints typically require an API key. Pass it via flag or env var:

If the model name differs from what can be derived from the endpoint URL (e.g. the hostname prefix), use --bi-encoder-model-name to set it explicitly.

Reusable, secret-free environment templates are provided for the recommended GPU configuration and the OpenAI cloud configuration:

• examples/env/best-quality.env.example
• examples/env/openai-cloud.env.example

Each file is a tracked template, not a runtime configuration. Copy the one you want to the ignored .env file at the repository root—or merge its settings into an existing .env—replace the placeholders there, and run the commented command with uv --env-file .env. Never commit the local .env containing real API keys.

Best quality (Qwen3 + GTE, both on GPU cluster):

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --base-url https://llm-serving.example.com/v1 --model my-model \
  --api-key $LLM_API_KEY \
  --bi-encoder-model https://qwen3-embedding-serving.example.com/v1 \
  --bi-encoder-model-name Qwen3-Embedding-4B \
  --bi-encoder-api-key $EMBEDDING_API_KEY \
  --cross-encoder-model https://gte-reranker-serving.example.com/v1/score \
  --query-instruction "Instruct: Given a text passage from an AI governance policy document, retrieve AI risk descriptions that are relevant to the concepts, requirements, or concerns discussed in the passage\nQuery: " \
  --expand-siblings

Good quality (EmbeddingGemma + GTE, remote):

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --base-url https://llm-serving.example.com/v1 --model my-model \
  --api-key $LLM_API_KEY \
  --bi-encoder-model https://embeddinggemma-serving.example.com/v1 \
  --bi-encoder-api-key $EMBEDDING_API_KEY \
  --cross-encoder-model https://gte-reranker-serving.example.com/v1/score

Local with GTE reranker (bi-encoder local, cross-encoder local — needs GPU for GTE):

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --base-url http://localhost:8000/v1 --model my-model \
  --cross-encoder-model Alibaba-NLP/gte-reranker-modernbert-base

Local defaults (no GPU needed, models downloaded automatically):

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --base-url http://localhost:8000/v1 --model my-model

IR-only (no LLM needed — useful for quick evaluation):

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --no-judge --no-grounding

OpenAI cloud model (GPT-6 Luna + OpenAI embeddings, no cross-encoder):

export OPENAI_API_KEY="..."
export NEXUS_BASE_DIR="/path/to/ai-atlas-nexus"
export POLICY_MAPPER_MAX_OUTPUT_TOKENS=8192
export POLICY_MAPPER_OUTPUT_TOKEN_PARAMETER=max_completion_tokens
# Optional: default, flex, fast, or priority. Omit to use the project setting.
export POLICY_MAPPER_SERVICE_TIER=default

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir "$NEXUS_BASE_DIR" \
  --base-url https://api.openai.com/v1 \
  --model gpt-6-luna \
  --api-key "$OPENAI_API_KEY" \
  --temperature 1.0 \
  --bi-encoder-model https://api.openai.com/v1 \
  --bi-encoder-model-name text-embedding-3-large \
  --bi-encoder-api-key "$OPENAI_API_KEY" \
  --no-cross-encoder \
  --query-instruction ""

This project currently uses OpenAI's Chat Completions integration for the LLM. The explicit output-token setting selects max_completion_tokens, while --query-instruction "" disables the taxonomy-specific embedding instruction. See the official GPT-6 Luna model documentation for model and endpoint details.

Requires Python 3.11+ and uv.

uv sync

You need a local clone of ai-atlas-nexus — set its path via NEXUS_BASE_DIR env var or --nexus-base-dir flag.

uv run asago-policy-mapper extract policy.pdf -o output/ \
  --base-url http://localhost:8000/v1 \
  --model my-model \
  --nexus-base-dir /path/to/ai-atlas-nexus

Outputs risk-extraction.json and risk-extraction.html report. Use --output-format yaml to get risk-extraction.yaml instead, or --output-format both for both.

uv run asago-policy-mapper eval output/ -g evals/ground_truth/policy-name.yaml

uv run python run_extract_battery.py batteries/risk-selected.yaml \
  --base-url <base-url> --model <model> -j 6

Runs extraction + eval across all policies in the battery config, generates per-run reports and a summary with per-taxonomy heatmaps.

# Battery with MLflow tracking disabled
uv run python run_extract_battery.py batteries/risk-selected.yaml \
  --base-url <url> --model <model> --no-mlflow

# Battery with custom MLflow experiment name
uv run python run_extract_battery.py batteries/risk-selected.yaml \
  --base-url <url> --model <model> --mlflow-experiment my-experiment

# IR-only mode (no LLM judge/grounding, no --base-url/--model needed)
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus --no-judge --no-grounding
uv run python run_extract_battery.py batteries/risk-selected.yaml \
  --no-judge --no-grounding

# Judge only, no grounding (test judge contribution in isolation)
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus --no-grounding \
  --base-url <url> --model <model>

# Skip causal synthesis (use static YAML chains)
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus --no-causal-synthesis \
  --base-url <url> --model <model>

# Smaller chunks with larger judge context window
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --chunk-max-tokens 256 --judge-context-tokens 512 \
  --base-url <url> --model <model>

# Disable LLM query generation (use raw chunk text for retrieval)
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --nexus-base-dir /path/to/ai-atlas-nexus --no-query-gen \
  --base-url <url> --model <model>

# Rebuild mitigation index (after data file changes)
python scripts/build_mitigation_index.py

You can provide your own risk taxonomy to extract against, alongside or instead of the built-in Nexus taxonomies. Custom risks go through the same retrieval, judging, and grounding pipeline as built-in risks.

# Required
taxonomy:
  id: my-taxonomy-id          # unique identifier, appears in extraction results
  name: My Taxonomy Name      # human-readable name
  description: "Optional description of the taxonomy"

# Required: at least one risk
risks:
  - id: my-risk-001           # unique risk identifier
    name: Risk Name            # human-readable name
    description: >-            # description used for retrieval (min 10 chars recommended)
      Full description of what this risk is about. The more detailed
      and specific the description, the better retrieval quality.
    concern: >-                # optional: additional context
      Why this risk matters and what harm it causes.
    risk_type: technical       # optional: "technical", "governance", "non-technical"

Required fields per risk: id, name, description. Optional fields: concern, risk_type.

See examples/custom-taxonomy-example.yaml for a complete example with 5 risks.

# Extract with a custom taxonomy alongside Nexus risks
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --base-url http://localhost:8000/v1 \
  --model my-model \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --custom-taxonomy my-risks.yaml

# Multiple custom taxonomies
uv run asago-policy-mapper extract policy.pdf -o output/ \
  --base-url http://localhost:8000/v1 \
  --model my-model \
  --nexus-base-dir /path/to/ai-atlas-nexus \
  --custom-taxonomy risks-a.yaml \
  --custom-taxonomy risks-b.yaml

Custom risks appear in extraction results with their taxonomy ID (e.g., taxonomy: "my-taxonomy-id"), grouped alongside built-in risks in the output report.

The loader validates your YAML before extraction starts:

• Missing or empty required fields produce clear error messages naming the specific risk
• Duplicate risk IDs within the file are rejected
• Very short descriptions (under 10 chars) trigger a warning
• Risk IDs using built-in prefixes (atlas-, nist-, owasp-, etc.) trigger a collision warning

• Custom risks do not have OWASP or NIST cross-taxonomy mappings
• Custom risks do not have mitigations or related actions
• Cross-taxonomy mapping generation for custom risks is not yet supported

# Unit tests (fast, no external dependencies)
uv run pytest tests/ -rs -m "not slow"

# Slow tests (loads embedding models)
uv run pytest tests/ -rs -m slow

# LLM integration tests (requires a local LLM server)
uv run pytest tests/ -rs --test-llm -m llm -v -s --tb=short -W ignore::DeprecationWarning

Integration tests marked @pytest.mark.llm exercise the full extraction pipeline against a real OpenAI-compatible LLM server. They validate pipeline mechanics (structured output, retries, token tracking, multi-stage orchestration) — not output quality.

With Ollama:

ollama pull gemma3:1b
uv run pytest tests/ -rs --test-llm -m llm -v -s --tb=short -W ignore::DeprecationWarning

Configuration:

Tests skip gracefully when no server is available.

Apache License 2.0