modules/cypher_graph_rag.py

"""CYPHER V12 M20 — Vector Graph RAG Engine.

Inspired by 2026 SOTA "Vector Graph RAG" paper (zilliztech) achieving
87.8% Recall@5 on multi-hop QA. Replaces flat vector retrieval with:

  1. Vector retrieval (BGE-M3 or all-MiniLM) over existing banks
  2. Knowledge Graph expansion (M28 CVE↔ATT&CK)
  3. Single reranking pass merging both signals
  4. Multi-hop: starting from prompt entities, expand graph 1-2 hops,
     re-query banks for each expanded entity, merge results

A2RAG (adaptive-agentic): controller decides if more evidence is needed,
otherwise stops early.

Used by cypher_bridge_v12 augment_prompt() as a richer alternative to
flat per-bank retrieval.
"""
from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import Any, Callable

logger = logging.getLogger(__name__)


@dataclass
class RagHit:
    text: str
    source: str  # bank name or "kg"
    score: float
    metadata: dict = field(default_factory=dict)


@dataclass
class GraphRagResult:
    hits: list[RagHit]
    context_string: str
    sources_used: list[str]
    multi_hop_used: bool
    iterations: int


class GraphRAGEngine:
    """Vector + Graph RAG orchestration.

    Args:
        vector_search_fn(query, k) -> list[RagHit]  (queries vector banks)
        kg: CveAttackKG instance (M28)
        evidence_threshold: when total scores >= this, stop early (A2RAG style)
    """

    def __init__(
        self,
        vector_search_fn: Callable[[str, int], list[RagHit]] | None = None,
        kg=None,
        evidence_threshold: float = 4.0,
        max_iterations: int = 2,
        k_per_query: int = 3,
    ):
        self.vector_search_fn = vector_search_fn
        self.kg = kg
        self.evidence_threshold = evidence_threshold
        self.max_iterations = max_iterations
        self.k_per_query = k_per_query

    @staticmethod
    def extract_entities(text: str) -> dict:
        cves = re.findall(r"CVE-\d{4}-\d{4,7}", text, re.IGNORECASE)
        tids = re.findall(r"\bT\d{4}(?:\.\d{3})?\b", text, re.IGNORECASE)
        vendors = []
        for v in ("Apache", "Microsoft", "Palo Alto", "Atlassian", "OpenSSL", "Fortinet", "Cisco"):
            if v.lower() in text.lower():
                vendors.append(v)
        return {
            "cves": [c.upper() for c in cves],
            "tids": [t.upper() for t in tids],
            "vendors": vendors,
        }

    def _initial_retrieval(self, query: str) -> list[RagHit]:
        hits: list[RagHit] = []
        if self.vector_search_fn:
            try:
                vec_hits = self.vector_search_fn(query, self.k_per_query)
                hits.extend(vec_hits or [])
            except Exception as e:
                logger.warning(f"vector_search fail: {e}")
        return hits

    def _kg_expansion(self, entities: dict) -> list[RagHit]:
        if self.kg is None:
            return []
        hits: list[RagHit] = []
        for cve in entities["cves"][:3]:
            ctx = self.kg.cve_full_context(cve)
            if ctx:
                # Build text representation
                text = (
                    f"{cve}: vendor={ctx.get('vendor')} cvss={ctx.get('cvss')} "
                    f"ransomware={ctx.get('ransomware')}"
                )
                tids = [n["node"] for n in ctx.get("tid", [])[:3]]
                if tids:
                    text += f" tids={tids}"
                mits = [n["node"] for n in ctx.get("mitigation", [])[:3]]
                if mits:
                    text += f" mitigations={mits}"
                hits.append(RagHit(text=text, source="kg_cve", score=2.5,
                                   metadata={"cve_id": cve}))
        for tid in entities["tids"][:3]:
            related = self.kg.tid_to_cves(tid)
            if related:
                hits.append(RagHit(
                    text=f"{tid} (MITRE) — observed in CVEs: {related[:3]}",
                    source="kg_tid",
                    score=2.0,
                    metadata={"tid": tid},
                ))
        for vendor in entities["vendors"][:2]:
            cves_v = self.kg.cves_by_vendor(vendor)
            if cves_v:
                hits.append(RagHit(
                    text=f"{vendor}: known critical CVEs {cves_v[:3]}",
                    source="kg_vendor",
                    score=1.5,
                    metadata={"vendor": vendor},
                ))
        return hits

    def _rerank(self, hits: list[RagHit], query: str) -> list[RagHit]:
        """Single reranking pass: boost hits matching query entities."""
        entities = self.extract_entities(query)
        kw_set = set()
        for k in entities["cves"] + entities["tids"] + entities["vendors"]:
            kw_set.add(k.lower())
        for h in hits:
            t = h.text.lower()
            for kw in kw_set:
                if kw.lower() in t:
                    h.score += 0.5
        # Source diversity: penalize duplicates of same source
        source_seen: dict[str, int] = {}
        for h in hits:
            source_seen[h.source] = source_seen.get(h.source, 0) + 1
            if source_seen[h.source] > 2:
                h.score -= 0.3
        return sorted(hits, key=lambda h: h.score, reverse=True)

    def _multi_hop_expand(self, query: str, current_hits: list[RagHit]) -> list[RagHit]:
        """For each top hit, extract entities and re-query."""
        if not self.vector_search_fn:
            return []
        expanded: list[RagHit] = []
        for h in current_hits[:3]:  # top 3 only
            new_entities = self.extract_entities(h.text)
            for cve in new_entities["cves"]:
                if cve.upper() not in query.upper():  # don't re-query already-known
                    try:
                        sub_hits = self.vector_search_fn(cve, 2)
                        for sh in sub_hits or []:
                            sh.score *= 0.7  # discount multi-hop
                            sh.metadata["multi_hop_via"] = h.source
                            expanded.append(sh)
                    except Exception:
                        continue
        return expanded

    def query(self, prompt: str, category: str | None = None) -> GraphRagResult:
        entities = self.extract_entities(prompt)
        all_hits: list[RagHit] = []
        sources_used: set[str] = set()
        multi_hop = False
        iters = 0

        # 1. Initial vector retrieval
        v_hits = self._initial_retrieval(prompt)
        all_hits.extend(v_hits)
        for h in v_hits:
            sources_used.add(h.source)

        # 2. KG expansion
        kg_hits = self._kg_expansion(entities)
        all_hits.extend(kg_hits)
        for h in kg_hits:
            sources_used.add(h.source)

        # 3. Rerank
        all_hits = self._rerank(all_hits, prompt)
        iters = 1

        # 4. Multi-hop if evidence insufficient (A2RAG)
        total_score = sum(h.score for h in all_hits[:5])
        if total_score < self.evidence_threshold and iters < self.max_iterations:
            expanded = self._multi_hop_expand(prompt, all_hits)
            if expanded:
                multi_hop = True
                all_hits.extend(expanded)
                all_hits = self._rerank(all_hits, prompt)
                iters += 1

        # 5. Build context
        top_hits = all_hits[: max(3, self.k_per_query)]
        ctx_parts: list[str] = []
        for h in top_hits:
            ctx_parts.append(f"[{h.source}] {h.text[:200]}")
        context_string = " | ".join(ctx_parts)
        if context_string:
            context_string = f"[GRAPH_RAG: {context_string}]"

        return GraphRagResult(
            hits=top_hits,
            context_string=context_string[:800],
            sources_used=sorted(sources_used),
            multi_hop_used=multi_hop,
            iterations=iters,
        )


__all__ = ["GraphRAGEngine", "RagHit", "GraphRagResult"]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M20 cypher_graph_rag SMOKE ===")

    # Mock vector search
    mock_corpus = {
        "log4shell jndi rce": "Log4Shell CVE-2021-44228 exploits JNDI lookups for RCE in Apache Log4j2.",
        "palo alto pan-os rce": "CVE-2024-3400 affects PAN-OS unauthenticated RCE.",
        "mitre att&ck initial access": "T1190 Exploit Public-Facing Application is common initial access vector.",
        "powershell encoded command": "T1059.001 PowerShell encoded commands often used by adversaries.",
        "lateral movement smb": "T1021.002 SMB lateral movement detected via psexec activity.",
    }
    def mock_vec(query, k):
        q = query.lower()
        scored = []
        for key, text in mock_corpus.items():
            overlap = sum(1 for tok in q.split() if tok in key.split())
            score = 1.5 + 0.3 * overlap
            scored.append(RagHit(text=text, source="MITRE", score=score))
        scored.sort(key=lambda h: h.score, reverse=True)
        return scored[:k]

    from cypher_kg_cve_attack import CveAttackKG
    kg = CveAttackKG()

    engine = GraphRAGEngine(
        vector_search_fn=mock_vec,
        kg=kg,
        evidence_threshold=5.0,
        k_per_query=3,
    )

    # Test 1: Simple CVE
    r = engine.query("What is CVE-2021-44228?", category="CYBERSEC")
    print(f"\n--- CVE-2021-44228 ---")
    print(f"Hits: {len(r.hits)}, sources: {r.sources_used}, multi_hop: {r.multi_hop_used}, iters: {r.iterations}")
    print(f"Top hit: [{r.hits[0].source}] {r.hits[0].text[:120]}")
    print(f"Context: {r.context_string[:300]}")

    # Test 2: Multi-entity
    r2 = engine.query("Analyze CVE-2021-44228 chain to T1190 with Apache mitigations", category="CYBERSEC")
    print(f"\n--- Multi-entity ---")
    print(f"Hits: {len(r2.hits)}, sources: {r2.sources_used}, multi_hop: {r2.multi_hop_used}")
    print(f"Context: {r2.context_string[:400]}")

    print("\n=== SMOKE PASS ===")