"""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 ===")