""" Risk knowledge base — FAISS index + RAG retrieval (Tasks 2.6.4, 2.6.5). Builds a FAISS index over the markdown reference docs in ``knowledge_base/`` (severity rubric, NIST CSF summary, US + international regulatory references), chunked to ~200 words and embedded with ``all-MiniLM-L6-v2``. ``retrieve_severity_context`` queries the index with a signal and returns the top-k most relevant chunks so the Risk Analysis Agent (Task 2.7) can ground its severity judgments. The embedder is injectable (shared with ``deduplication``) so tests can run with a deterministic fake and no model download. """ from __future__ import annotations import json import logging import re from pathlib import Path from typing import Optional, Sequence import numpy as np from src.analysis.deduplication import Embedder, get_default_embedder logger = logging.getLogger(__name__) # knowledge_base/ sits at the project root (two levels up from this file: src/analysis/). _PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent KB_DIR = _PROJECT_ROOT / "knowledge_base" KB_FILES = ( "severity_rubric.md", "nist_csf_summary.md", "regulatory_reference.md", "regulatory_reference_intl.md", ) INDEX_PATH = KB_DIR / "severity.index" CHUNKS_PATH = KB_DIR / "severity_chunks.json" CHUNK_TARGET_WORDS = 200 # In-memory cache keyed by index path so repeated retrievals (and tests using # different paths) don't reload or collide. _cache: dict[str, tuple] = {} # ── Chunking ────────────────────────────────────────────────────────────────── def chunk_markdown(text: str, *, source: str = "", target_words: int = CHUNK_TARGET_WORDS) -> list[str]: """Split markdown into ~target_words chunks, keeping the nearest heading as a prefix so each chunk carries its context (e.g. 'CYBERSECURITY'). Words from each heading section accumulate and are emitted in fixed target_words windows, so even a single very long paragraph is split. """ heading = "" chunks: list[str] = [] words: list[str] = [] def flush() -> None: nonlocal words prefix = f"[{source}] " if source else "" ctx = f"{heading}: " if heading else "" for i in range(0, len(words), target_words): body = " ".join(words[i : i + target_words]).strip() if body: chunks.append(f"{prefix}{ctx}{body}") words = [] for raw in text.splitlines(): line = raw.strip() m = re.match(r"^(#{1,6})\s+(.*)$", line) if m: flush() heading = m.group(2).strip() continue if not line or re.fullmatch(r"-{3,}", line): # blank line or horizontal rule continue words.extend(line.split()) flush() return chunks def load_kb_chunks(doc_paths: Optional[Sequence[Path]] = None) -> list[str]: """Read and chunk all knowledge-base documents.""" paths = [Path(p) for p in (doc_paths or [KB_DIR / f for f in KB_FILES])] chunks: list[str] = [] for path in paths: if not path.exists(): logger.warning("knowledge-base doc missing: %s", path) continue text = path.read_text(encoding="utf-8") chunks.extend(chunk_markdown(text, source=path.stem)) return chunks # ── Index build / load ──────────────────────────────────────────────────────── def build_index( *, embedder: Optional[Embedder] = None, doc_paths: Optional[Sequence[Path]] = None, index_path: Path = INDEX_PATH, chunks_path: Path = CHUNKS_PATH, ) -> int: """Build the FAISS index over the knowledge base and write it to disk. Returns the number of chunks indexed. Uses inner-product over L2-normalised embeddings, i.e. cosine similarity. """ import faiss # deferred — heavy native dependency embedder = embedder or get_default_embedder() chunks = load_kb_chunks(doc_paths) if not chunks: raise RuntimeError(f"No knowledge-base chunks found (looked in {KB_DIR}).") emb = np.asarray(embedder(chunks), dtype="float32") faiss.normalize_L2(emb) index = faiss.IndexFlatIP(emb.shape[1]) index.add(emb) index_path.parent.mkdir(parents=True, exist_ok=True) faiss.write_index(index, str(index_path)) chunks_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8") _cache.pop(str(index_path), None) logger.info("Built knowledge-base index: %d chunks -> %s", len(chunks), index_path) return len(chunks) def _load_index(index_path: Path, chunks_path: Path): key = str(index_path) if key in _cache: return _cache[key] import faiss if not index_path.exists() or not chunks_path.exists(): logger.info("Knowledge-base index missing — building it now.") build_index(index_path=index_path, chunks_path=chunks_path) index = faiss.read_index(str(index_path)) chunks = json.loads(chunks_path.read_text(encoding="utf-8")) _cache[key] = (index, chunks) return index, chunks # ── Retrieval (Task 2.6.5) ──────────────────────────────────────────────────── def _signal_query(signal) -> str: """Build a retrieval query from a RiskSignal (category + subcategory + text).""" cat = getattr(getattr(signal, "risk_category", None), "value", "") or "" sub = getattr(signal, "risk_subcategory", "") or "" text = getattr(signal, "text", "") or "" return " ".join(p for p in (cat, sub, text) if p).strip() async def retrieve_severity_context( signal, k: int = 3, *, embedder: Optional[Embedder] = None, index_path: Path = INDEX_PATH, chunks_path: Path = CHUNKS_PATH, ) -> list[str]: """Return the top-k knowledge-base chunks most relevant to *signal*. *signal* is a ``RiskSignal`` (or anything with risk_category / risk_subcategory / text). The index is built lazily on first use if absent. """ import asyncio import faiss query = _signal_query(signal) if not query: return [] def _search() -> list[str]: index, chunks = _load_index(index_path, chunks_path) emb_fn = embedder or get_default_embedder() q = np.asarray(emb_fn([query]), dtype="float32") faiss.normalize_L2(q) _scores, idxs = index.search(q, min(k, len(chunks))) return [chunks[i] for i in idxs[0] if i != -1] return await asyncio.to_thread(_search) __all__ = [ "KB_DIR", "KB_FILES", "INDEX_PATH", "CHUNKS_PATH", "chunk_markdown", "load_kb_chunks", "build_index", "retrieve_severity_context", ] if __name__ == "__main__": # `python -m src.analysis.knowledge_base` builds the index logging.basicConfig(level=logging.INFO) n = build_index() print(f"Indexed {n} knowledge-base chunks at {INDEX_PATH}")