icdev 0.0.3__py3-none-any.whl

tools/knowledge_graph/graph_rag.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python3
+# CUI // SP-CTI
+"""GraphRAG Integration — Graph-augmented retrieval for ICDEV RAG pipeline.
+
+Overlays user queries onto knowledge graphs to retrieve relational context
+that standard vector/keyword RAG misses. Follows InfraNodus GraphRAG pattern:
+  1. Convert query → mini-graph (co-occurrence)
+  2. Match query nodes to knowledge graph nodes
+  3. Follow edges to retrieve related concepts + their communities
+  4. Return enriched context with structural relationships
+
+Architecture decisions:
+  - D-KG-8: GraphRAG enhances (not replaces) existing corrective_rag.py
+  - D-KG-9: Query expansion via graph neighborhood traversal (BFS depth=2)
+  - D-KG-10: Scoring = 0.4 * relevance + 0.3 * centrality + 0.3 * proximity
+  - D-KARL-1: Per-query-type scoring profiles (compliance, exploratory, entity, synthesis)
+  - D-KARL-2: Self-directed context compression via scanner-tier LLM
+
+Usage:
+    python tools/knowledge_graph/graph_rag.py --query "zero trust architecture" --graph-id <id> --json
+    python tools/knowledge_graph/graph_rag.py --query "compliance gaps" --project-id sparkpilot --json
+    python tools/knowledge_graph/graph_rag.py --query "AC-2" --profile compliance --json
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+BASE_DIR = Path(__file__).resolve().parent.parent.parent
+sys.path.insert(0, str(BASE_DIR))
+
+from tools.knowledge_graph.text_network import (  # noqa: E402
+    tokenize, load_graph, list_graphs, _ensure_tables,
+)
+from tools.db.storage import get_connection  # noqa: E402
+
+# ── Default scoring weights ──────────────────────────────────────
+W_RELEVANCE = 0.4   # Direct match to query terms
+W_CENTRALITY = 0.3  # Betweenness centrality (importance in graph)
+W_PROXIMITY = 0.3   # Graph distance from query terms
+
+# ── KARL-inspired scoring profiles (D-KARL-1) ────────────────────
+_SCORING_PROFILES: dict[str, dict[str, float]] = {
+    "default": {"relevance": 0.4, "centrality": 0.3, "proximity": 0.3},
+    "compliance": {"relevance": 0.3, "centrality": 0.5, "proximity": 0.2},
+    "exploratory": {"relevance": 0.3, "centrality": 0.2, "proximity": 0.5},
+    "entity_search": {"relevance": 0.6, "centrality": 0.2, "proximity": 0.2},
+    "synthesis": {"relevance": 0.35, "centrality": 0.35, "proximity": 0.3},
+}
+
+# Heuristic patterns for auto-detecting query type
+_COMPLIANCE_RE = re.compile(
+    r"\b([A-Z]{2})-\d|NIST|FedRAMP|CMMC|STIG|ATO|cATO|POAM|SSP|OSCAL|compliance|control",
+    re.IGNORECASE,
+)
+_ENTITY_RE = re.compile(r'^"[^"]+"|^\S+$|^what is |^define ', re.IGNORECASE)
+
+
+def _load_scoring_profiles() -> dict[str, dict[str, float]]:
+    """Load scoring profiles from args/knowledge_graph_config.yaml."""
+    config_path = BASE_DIR / "args" / "knowledge_graph_config.yaml"
+    if config_path.exists():
+        try:
+            import yaml
+            with open(config_path) as f:
+                cfg = yaml.safe_load(f) or {}
+            profiles = cfg.get("graph_rag", {}).get("scoring_profiles", {})
+            result = dict(_SCORING_PROFILES)
+            for name, weights in profiles.items():
+                result[name] = {
+                    "relevance": weights.get("relevance_weight", 0.4),
+                    "centrality": weights.get("centrality_weight", 0.3),
+                    "proximity": weights.get("proximity_weight", 0.3),
+                }
+            return result
+        except (ImportError, Exception):
+            pass
+    return dict(_SCORING_PROFILES)
+
+
+def _load_compression_config() -> dict[str, Any]:
+    """Load compression settings from args/knowledge_graph_config.yaml."""
+    config_path = BASE_DIR / "args" / "knowledge_graph_config.yaml"
+    if config_path.exists():
+        try:
+            import yaml
+            with open(config_path) as f:
+                cfg = yaml.safe_load(f) or {}
+            return cfg.get("graph_rag", {}).get("compression", {
+                "enabled": True, "max_context_tokens": 2000,
+                "llm_function": "kg_compress",
+            })
+        except (ImportError, Exception):
+            pass
+    return {"enabled": True, "max_context_tokens": 2000, "llm_function": "kg_compress"}
+
+
+def detect_query_profile(query: str) -> str:
+    """Auto-detect the best scoring profile for a query (D-KARL-1).
+
+    Returns profile name: compliance, entity_search, exploratory, synthesis, or default.
+    """
+    if _COMPLIANCE_RE.search(query):
+        return "compliance"
+    if _ENTITY_RE.search(query):
+        return "entity_search"
+    # Multi-topic queries → synthesis
+    conjunctions = len(re.findall(r"\band\b|\bor\b|[;,]", query, re.IGNORECASE))
+    if conjunctions >= 2:
+        return "synthesis"
+    # Questions with "how", "why", "what if" → exploratory
+    if re.search(r"\bhow\b|\bwhy\b|\bwhat if\b|\bexplore\b|\bgap\b", query, re.IGNORECASE):
+        return "exploratory"
+    return "default"
+
+
+def get_scoring_weights(profile: str | None = None, query: str = "") -> tuple[float, float, float]:
+    """Get (relevance, centrality, proximity) weights for a profile.
+
+    If profile is None, auto-detects from query text.
+    """
+    profiles = _load_scoring_profiles()
+    if profile is None:
+        profile = detect_query_profile(query)
+    weights = profiles.get(profile, profiles["default"])
+    return weights["relevance"], weights["centrality"], weights["proximity"]
+
+
+def compress_context(context: str, query: str, *, max_tokens: int = 2000) -> str:
+    """KARL-inspired self-directed context compression (D-KARL-2).
+
+    Compresses verbose GraphRAG context into a tighter representation
+    using scanner-tier LLM (qwen3.5, zero Claude tokens). Falls back
+    to deterministic truncation if LLM unavailable.
+    """
+    # Rough token estimate: 1 token ≈ 4 chars
+    estimated_tokens = len(context) // 4
+    if estimated_tokens <= max_tokens:
+        return context
+
+    config = _load_compression_config()
+    if not config.get("enabled", True):
+        # Truncate deterministically
+        return context[: max_tokens * 4]
+
+    try:
+        from tools.llm.router import invoke
+        prompt = (
+            f"Compress this knowledge graph retrieval context for the query: \"{query}\"\n\n"
+            f"CONTEXT:\n{context}\n\n"
+            f"INSTRUCTIONS:\n"
+            f"- Keep all concept names, scores, and community IDs\n"
+            f"- Remove redundant neighbors and low-score entries\n"
+            f"- Preserve structural gaps (blind spots) — these are high value\n"
+            f"- Keep the most central and query-relevant concepts\n"
+            f"- Output a compressed version under {max_tokens} tokens\n"
+            f"- Do NOT add commentary — output ONLY the compressed context"
+        )
+        result = invoke(
+            prompt=prompt,
+            function_name=config.get("llm_function", "kg_compress"),
+            max_tokens=max_tokens,
+        )
+        if isinstance(result, dict):
+            compressed = result.get("content", result.get("text", ""))
+        else:
+            compressed = str(result)
+        if compressed and len(compressed) > 50:
+            return compressed
+    except Exception:
+        pass
+
+    # Deterministic fallback: keep first N lines
+    lines = context.split("\n")
+    truncated = []
+    char_count = 0
+    for line in lines:
+        if char_count + len(line) > max_tokens * 4:
+            break
+        truncated.append(line)
+        char_count += len(line) + 1
+    return "\n".join(truncated)
+
+
+def _graph_neighborhood(
+    adj: dict[str, dict[str, float]],
+    seed_nodes: set[str],
+    max_depth: int = 2,
+) -> dict[str, int]:
+    """BFS from seed nodes, return {node: distance}."""
+    visited: dict[str, int] = {n: 0 for n in seed_nodes if n in adj}
+    queue = list(visited.keys())
+
+    for depth in range(1, max_depth + 1):
+        next_queue = []
+        for node in queue:
+            for neighbor in adj.get(node, {}):
+                if neighbor not in visited:
+                    visited[neighbor] = depth
+                    next_queue.append(neighbor)
+        queue = next_queue
+
+    return visited
+
+
+def _build_adj_from_graph(graph_data: dict) -> dict[str, dict[str, float]]:
+    """Build adjacency dict from loaded graph data."""
+    adj: dict[str, dict[str, float]] = defaultdict(dict)
+    for edge in graph_data.get("edges", []):
+        s, t, w = edge["source"], edge["target"], edge.get("weight", 1.0)
+        adj[s][t] = w
+        adj[t][s] = w
+    return dict(adj)
+
+
+def graph_rag_retrieve(
+    query: str,
+    graph_data: dict,
+    *,
+    max_depth: int = 2,
+    top_k: int = 20,
+    profile: str | None = None,
+    compress: bool = True,
+) -> dict[str, Any]:
+    """Graph-augmented retrieval from a knowledge graph.
+
+    1. Tokenize query → seed nodes
+    2. BFS neighborhood expansion (depth=2)
+    3. Score nodes using profile-aware weights (D-KARL-1)
+    4. Return ranked concepts with community context
+    5. Optionally compress context for downstream consumption (D-KARL-2)
+
+    Args:
+        profile: Scoring profile name (compliance, exploratory, entity_search,
+                 synthesis). None = auto-detect from query text.
+        compress: If True, compress context via scanner-tier LLM when large.
+
+    Returns enriched context for downstream LLM consumption.
+    """
+    query_tokens = set(tokenize(query, lemmatize=True))
+    if not query_tokens:
+        return {"error": "No meaningful tokens in query"}
+
+    # Resolve scoring weights from profile (D-KARL-1)
+    detected_profile = profile or detect_query_profile(query)
+    w_rel, w_cen, w_prox = get_scoring_weights(detected_profile, query)
+
+    # Build lookup structures
+    node_map = {n["id"]: n for n in graph_data.get("nodes", [])}
+    adj = _build_adj_from_graph(graph_data)
+
+    # Find matching seed nodes
+    seed_nodes = query_tokens & set(node_map.keys())
+
+    # If no exact matches, try partial matching
+    if not seed_nodes:
+        for qt in query_tokens:
+            for node_id in node_map:
+                if qt in node_id or node_id in qt:
+                    seed_nodes.add(node_id)
+
+    if not seed_nodes:
+        return {
+            "query": query,
+            "query_tokens": list(query_tokens),
+            "matched_seeds": [],
+            "results": [],
+            "scoring_profile": detected_profile,
+            "context": "No matching concepts found in knowledge graph.",
+        }
+
+    # BFS neighborhood
+    distances = _graph_neighborhood(adj, seed_nodes, max_depth)
+
+    # Score all reachable nodes
+    max_centrality = max(
+        (n.get("centrality", 0) for n in node_map.values()), default=1.0
+    ) or 1.0
+
+    scored: list[dict[str, Any]] = []
+    for node_id, depth in distances.items():
+        node = node_map.get(node_id)
+        if not node:
+            continue
+
+        # Relevance: 1.0 for direct match, 0.5 for partial
+        relevance = 1.0 if node_id in query_tokens else 0.5
+        # Centrality: normalized
+        centrality = node.get("centrality", 0) / max_centrality
+        # Proximity: inverse of distance (closer = better)
+        proximity = 1.0 / (1 + depth)
+
+        score = (w_rel * relevance +
+                 w_cen * centrality +
+                 w_prox * proximity)
+
+        scored.append({
+            "concept": node_id,
+            "score": round(score, 4),
+            "relevance": round(relevance, 4),
+            "centrality_score": round(centrality, 4),
+            "proximity_score": round(proximity, 4),
+            "community": node.get("community", 0),
+            "occurrences": node.get("occurrences", 0),
+            "graph_distance": depth,
+            "neighbors": list(adj.get(node_id, {}).keys())[:10],
+        })
+
+    scored.sort(key=lambda x: -x["score"])
+    top_results = scored[:top_k]
+
+    # Build community context
+    communities_seen: dict[int, list[str]] = defaultdict(list)
+    for r in top_results:
+        communities_seen[r["community"]].append(r["concept"])
+
+    # Generate context string for LLM
+    context_parts = [f"Query: {query}", ""]
+    context_parts.append("Related concepts from knowledge graph:")
+    for r in top_results[:10]:
+        context_parts.append(
+            f"  - {r['concept']} (community {r['community']}, "
+            f"centrality={r['centrality_score']:.2f}, "
+            f"neighbors: {', '.join(r['neighbors'][:5])})"
+        )
+
+    # Add structural gap context
+    gaps = graph_data.get("structural_gaps", [])
+    relevant_gaps = [
+        g for g in gaps
+        if any(c in communities_seen for c in [g.get("community_a"), g.get("community_b")])
+    ]
+    if relevant_gaps:
+        context_parts.append("\nStructural gaps (potential blind spots):")
+        for g in relevant_gaps[:3]:
+            context_parts.append(
+                f"  Gap between community {g['community_a']} "
+                f"({', '.join(g.get('community_a_top_concepts', [])[:3])}) "
+                f"and community {g['community_b']} "
+                f"({', '.join(g.get('community_b_top_concepts', [])[:3])})"
+                f" — gap_score={g['gap_score']}"
+            )
+
+    raw_context = "\n".join(context_parts)
+
+    # KARL self-directed compression (D-KARL-2)
+    if compress:
+        config = _load_compression_config()
+        max_tokens = config.get("max_context_tokens", 2000)
+        final_context = compress_context(raw_context, query, max_tokens=max_tokens)
+        compressed = final_context != raw_context
+    else:
+        final_context = raw_context
+        compressed = False
+
+    return {
+        "query": query,
+        "query_tokens": list(query_tokens),
+        "matched_seeds": list(seed_nodes),
+        "results": top_results,
+        "scoring_profile": detected_profile,
+        "scoring_weights": {"relevance": w_rel, "centrality": w_cen, "proximity": w_prox},
+        "communities_involved": {
+            k: v for k, v in communities_seen.items()
+        },
+        "relevant_gaps": relevant_gaps[:5],
+        "context": final_context,
+        "context_compressed": compressed,
+        "classification": "CUI // SP-CTI",
+    }
+
+
+def graph_rag_query(
+    query: str,
+    *,
+    graph_id: str | None = None,
+    project_id: str = "sparkpilot",
+    max_depth: int = 2,
+    top_k: int = 20,
+    profile: str | None = None,
+    compress: bool = True,
+) -> dict[str, Any]:
+    """High-level GraphRAG query interface.
+
+    If graph_id given, queries that specific graph.
+    Otherwise, queries all graphs for the project and merges results.
+
+    Args:
+        profile: Scoring profile (compliance, exploratory, entity_search,
+                 synthesis). None = auto-detect from query.
+        compress: Compress context via scanner-tier LLM (D-KARL-2).
+    """
+    _ensure_tables()
+
+    if graph_id:
+        graph_data = load_graph(graph_id)
+        if not graph_data:
+            return {"error": f"Graph {graph_id} not found"}
+        result = graph_rag_retrieve(query, graph_data,
+                                    max_depth=max_depth, top_k=top_k,
+                                    profile=profile, compress=compress)
+        result["graph_id"] = graph_id
+        return result
+
+    # Query all project graphs and merge
+    graphs = list_graphs(project_id)
+    if not graphs:
+        return {"error": f"No knowledge graphs found for project {project_id}"}
+
+    detected_profile = profile or detect_query_profile(query)
+    all_results = []
+    for g in graphs[:10]:  # Limit to 10 most recent graphs
+        graph_data = load_graph(g["id"])
+        if not graph_data:
+            continue
+        result = graph_rag_retrieve(query, graph_data,
+                                    max_depth=max_depth, top_k=top_k,
+                                    profile=detected_profile, compress=False)
+        for r in result.get("results", []):
+            r["source_graph"] = g["id"]
+            r["source_graph_name"] = g["name"]
+        all_results.extend(result.get("results", []))
+
+    # De-duplicate and re-rank
+    seen: dict[str, dict] = {}
+    for r in all_results:
+        concept = r["concept"]
+        if concept not in seen or r["score"] > seen[concept]["score"]:
+            seen[concept] = r
+
+    merged = sorted(seen.values(), key=lambda x: -x["score"])[:top_k]
+
+    # Build merged context + compress (D-KARL-2)
+    context_parts = [f"Query: {query}", ""]
+    context_parts.append("Related concepts (merged from multiple graphs):")
+    for r in merged[:10]:
+        context_parts.append(
+            f"  - {r['concept']} (score={r['score']}, "
+            f"community={r.get('community', '?')}, "
+            f"graph={r.get('source_graph_name', '?')})"
+        )
+    raw_context = "\n".join(context_parts)
+
+    if compress:
+        config = _load_compression_config()
+        final_context = compress_context(
+            raw_context, query,
+            max_tokens=config.get("max_context_tokens", 2000),
+        )
+    else:
+        final_context = raw_context
+
+    return {
+        "query": query,
+        "graphs_searched": len(graphs),
+        "scoring_profile": detected_profile,
+        "results": merged,
+        "context": final_context,
+        "classification": "CUI // SP-CTI",
+    }
+
+
+# ── CLI ──────────────────────────────────────────────────────────
+def main():
+    parser = argparse.ArgumentParser(description="GraphRAG Knowledge Retrieval")
+    parser.add_argument("--query", required=True, help="Search query")
+    parser.add_argument("--graph-id", help="Specific graph to query")
+    parser.add_argument("--project-id", default="sparkpilot")
+    parser.add_argument("--max-depth", type=int, default=2)
+    parser.add_argument("--top-k", type=int, default=20)
+    parser.add_argument("--profile", choices=list(_SCORING_PROFILES.keys()),
+                        default=None,
+                        help="Scoring profile (auto-detected if omitted)")
+    parser.add_argument("--no-compress", action="store_true",
+                        help="Disable context compression")
+    parser.add_argument("--json", action="store_true")
+    args = parser.parse_args()
+
+    result = graph_rag_query(
+        args.query,
+        graph_id=args.graph_id,
+        project_id=args.project_id,
+        max_depth=args.max_depth,
+        top_k=args.top_k,
+        profile=args.profile,
+        compress=not args.no_compress,
+    )
+
+    print(json.dumps(result, indent=2, default=str) if args.json else str(result))
+
+
+if __name__ == "__main__":
+    main()