merge: chore/cleanup-remove-bloat-and-secrets into main

app/knowledge_graph.py

@@ -0,0 +1,772 @@
+#!/usr/bin/env python3
+"""
+Knowledge Graph — Wallet → Token → Scam Relationship Engine
+=============================================================
+Builds and queries a knowledge graph from RAG document metadata.
+
+Nodes:  wallet, token, contract, scam_pattern, entity, chain
+Edges:  holds_token, deployed_by, involved_in, similar_to, flags, transfers_to
+
+Powered by Redis for real-time graph storage + traversal.
+Used by:
+  - three_pillar_search() as enhanced entity lookup (Pillar 3+)
+  - Relationship queries ("who is connected to this wallet?")
+  - Forensic tracing (find paths between addresses)
+
+Architecture:
+  - Redis graph stored as adjacency lists (sorted sets with typed edges)
+  - Edge types: HOLDS, DEPLOYED, INVOLVED_IN, SIMILAR_TO, FLAGS, TRANSFERS
+  - Each edge has a weight (confidence 0-1) and metadata
+  - Auto-built from document metadata during ingestion
+  - Supports: neighbor queries, pathfinding (BFS up to depth 4), subgraph extraction
+
+Key naming:
+  kg:node:{type}:{id}           -> hash  {label, metadata_json}
+  kg:edges:{type}:{id}           -> sorted set  member=edge_key,  score=weight
+  kg:edge:{from_type}:{from_id}:{rel}:{to_type}:{to_id}  -> hash  {weight, metadata_json, created_at}
+"""
+
+import contextlib
+import json
+import logging
+import os
+import time
+from typing import Any
+
+import redis.asyncio as aioredis
+
+logger = logging.getLogger(__name__)
+
+REDIS_HOST = os.getenv("REDIS_HOST", "rmi-redis")
+REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
+REDIS_PASSWORD = os.environ.get("REDIS_PASSWORD", "")
+
+# ── Edge Types ────────────────────────────────────────────────────
+
+
+class EdgeType:
+    HOLDS = "holds"  # wallet holds token
+    DEPLOYED = "deployed"  # wallet deployed contract/token
+    INVOLVED_IN = "involved"  # wallet involved in scam/event
+    SIMILAR_TO = "similar"  # token/wallet similar to another
+    FLAGS = "flags"  # entity flags another as suspicious
+    TRANSFERS = "transfers"  # wallet transfers to another wallet
+    OWNS = "owns"  # wallet owns NFT/token
+    MIGRATED = "migrated"  # token migrated to new contract
+    ASSOCIATED = "associated"  # generic association
+
+
+# ── Node Types ────────────────────────────────────────────────────
+
+
+class NodeType:
+    WALLET = "wallet"
+    TOKEN = "token"
+    CONTRACT = "contract"
+    SCAM = "scam"
+    ENTITY = "entity"  # named entity (protocol, person, org)
+    CHAIN = "chain"
+    ADDRESS = "address"  # generic blockchain address
+
+
+# ── Knowledge Graph Engine ─────────────────────────────────────────
+
+
+class KnowledgeGraph:
+    """
+    Redis-backed knowledge graph for crypto entity relationships.
+
+    Stores nodes and edges in Redis for fast graph traversal.
+    Auto-built from RAG document metadata during ingestion.
+    """
+
+    def __init__(self, redis: aioredis.Redis = None):
+        self._redis = redis
+        self._node_cache: dict[str, dict] = {}
+
+    async def _get_redis(self) -> aioredis.Redis:
+        if self._redis:
+            return self._redis
+        self._redis = aioredis.Redis(
+            host=REDIS_HOST,
+            port=REDIS_PORT,
+            password=REDIS_PASSWORD or None,
+            db=0,
+            decode_responses=True,
+        )
+        return self._redis
+
+    # ── Node Operations ───────────────────────────────────────────
+
+    async def add_node(
+        self,
+        node_type: str,
+        node_id: str,
+        label: str = "",
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        """Add or update a node in the knowledge graph."""
+        r = await self._get_redis()
+        key = f"kg:node:{node_type}:{node_id}"
+        data = {
+            "label": label or node_id,
+            "type": node_type,
+            "metadata": json.dumps(metadata or {}),
+            "updated_at": str(int(time.time())),
+        }
+        await r.hset(key, mapping=data)
+        self._node_cache[f"{node_type}:{node_id}"] = data
+        return key
+
+    async def get_node(self, node_type: str, node_id: str) -> dict | None:
+        """Get a node by type and ID."""
+        r = await self._get_redis()
+        key = f"kg:node:{node_type}:{node_id}"
+        data = await r.hgetall(key)
+        if not data:
+            return None
+        result = dict(data)
+        if "metadata" in result and isinstance(result["metadata"], str):
+            with contextlib.suppress(json.JSONDecodeError):
+                result["metadata"] = json.loads(result["metadata"])
+        return result
+
+    # ── Edge Operations ────────────────────────────────────────────
+
+    async def add_edge(
+        self,
+        from_type: str,
+        from_id: str,
+        relation: str,
+        to_type: str,
+        to_id: str,
+        weight: float = 1.0,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:
+        """
+        Add a directed edge between two nodes.
+        Weight: 0.0-1.0 confidence/strength score.
+        """
+        r = await self._get_redis()
+        edge_key = f"kg:edge:{from_type}:{from_id}:{relation}:{to_type}:{to_id}"
+
+        edge_data = {
+            "from_type": from_type,
+            "from_id": from_id,
+            "relation": relation,
+            "to_type": to_type,
+            "to_id": to_id,
+            "weight": str(weight),
+            "metadata": json.dumps(metadata or {}),
+            "created_at": str(int(time.time())),
+        }
+        await r.hset(edge_key, mapping=edge_data)
+
+        # Add to from-node's edge index (sorted set, score=weight)
+        from_edge_key = f"kg:edges:{from_type}:{from_id}"
+        await r.zadd(from_edge_key, {edge_key: weight})
+
+        # Add reverse reference for to-node (for incoming edge queries)
+        to_edge_key = f"kg:edges_in:{to_type}:{to_id}"
+        await r.zadd(to_edge_key, {edge_key: weight})
+
+        return edge_key
+
+    async def get_neighbors(
+        self,
+        node_type: str,
+        node_id: str,
+        relation: str | None = None,
+        direction: str = "outgoing",
+        min_weight: float = 0.0,
+        limit: int = 50,
+    ) -> list[dict]:
+        """
+        Get neighboring nodes of a given node.
+
+        Args:
+            direction: 'outgoing' (default), 'incoming', or 'both'
+            relation: filter by edge type (None = all)
+            min_weight: minimum edge weight threshold
+            limit: max results
+        """
+        r = await self._get_redis()
+        results = []
+
+        # Build list of direction prefixes to search (fixes Python ternary precedence bug)
+        dirs = []
+        if direction in ("outgoing", "both"):
+            dirs.append(("kg:edges:", "outgoing"))
+        if direction in ("incoming", "both"):
+            dirs.append(("kg:edges_in:", "incoming"))
+        for dir_key in dirs:
+            prefix = dir_key[0]
+            idx_key = f"{prefix}{node_type}:{node_id}"
+            edge_keys = await r.zrevrangebyscore(idx_key, "+inf", min_weight)
+            for ek in edge_keys[:limit]:
+                edge_data = await r.hgetall(ek)
+                if not edge_data:
+                    continue
+                if relation and edge_data.get("relation") != relation:
+                    continue
+                weight = float(edge_data.get("weight", 0))
+                if weight < min_weight:
+                    continue
+                # Determine the neighbor
+                if dir_key[1] == "outgoing":
+                    neighbor_type = edge_data.get("to_type", "")
+                    neighbor_id = edge_data.get("to_id", "")
+                else:
+                    neighbor_type = edge_data.get("from_type", "")
+                    neighbor_id = edge_data.get("from_id", "")
+
+                neighbor_node = await self.get_node(neighbor_type, neighbor_id)
+                results.append(
+                    {
+                        "node_type": neighbor_type,
+                        "node_id": neighbor_id,
+                        "label": neighbor_node.get("label", neighbor_id) if neighbor_node else neighbor_id,
+                        "relation": edge_data.get("relation", ""),
+                        "weight": weight,
+                        "direction": dir_key[1],
+                        "metadata": neighbor_node.get("metadata", {}) if neighbor_node else {},
+                        "edge_metadata": json.loads(edge_data.get("metadata", "{}")),
+                    }
+                )
+                if len(results) >= limit:
+                    return results
+
+        return results
+
+    async def find_paths(
+        self,
+        from_type: str,
+        from_id: str,
+        to_type: str,
+        to_id: str,
+        max_depth: int = 4,
+        max_paths: int = 5,
+    ) -> list[list[dict]]:
+        """
+        BFS pathfinding between two nodes.
+        Returns up to max_paths shortest paths.
+        """
+        r = await self._get_redis()
+        paths = []
+
+        # BFS
+        queue = [(from_type, from_id, [])]
+        visited = {(from_type, from_id)}
+
+        while queue and len(paths) < max_paths:
+            curr_type, curr_id, path_so_far = queue.pop(0)
+
+            if len(path_so_far) >= max_depth:
+                continue
+
+            # Get all edges from this node
+            idx_key = f"kg:edges:{curr_type}:{curr_id}"
+            edge_keys = await r.zrevrangebyscore(idx_key, "+inf", "0")
+
+            for ek in edge_keys:
+                edge_data = await r.hgetall(ek)
+                if not edge_data:
+                    continue
+
+                neighbor_type = edge_data.get("to_type", "")
+                neighbor_id = edge_data.get("to_id", "")
+
+                if (neighbor_type, neighbor_id) in visited:
+                    continue
+
+                step = {
+                    "from_type": curr_type,
+                    "from_id": curr_id,
+                    "relation": edge_data.get("relation", ""),
+                    "to_type": neighbor_type,
+                    "to_id": neighbor_id,
+                    "weight": float(edge_data.get("weight", 0)),
+                }
+
+                new_path = [*path_so_far, step]
+
+                # Check if we reached the target
+                if neighbor_type == to_type and neighbor_id == to_id:
+                    paths.append(new_path)
+                    continue
+
+                visited.add((neighbor_type, neighbor_id))
+                queue.append((neighbor_type, neighbor_id, new_path))
+
+        return paths
+
+    async def get_subgraph(
+        self,
+        node_type: str,
+        node_id: str,
+        depth: int = 2,
+        max_nodes: int = 100,
+    ) -> dict[str, Any]:
+        """
+        Extract a subgraph around a node for visualization/analysis.
+        Returns nodes and edges for graph rendering.
+        """
+        nodes = {}
+        edges = []
+        visited = set()
+        queue = [(node_type, node_id, 0)]
+
+        while queue and len(nodes) < max_nodes:
+            curr_type, curr_id, curr_depth = queue.pop(0)
+
+            if (curr_type, curr_id) in visited:
+                continue
+            if curr_depth > depth:
+                continue
+
+            visited.add((curr_type, curr_id))
+            node_data = await self.get_node(curr_type, curr_id)
+            nodes[f"{curr_type}:{curr_id}"] = {
+                "type": curr_type,
+                "id": curr_id,
+                "label": node_data.get("label", curr_id) if node_data else curr_id,
+                "metadata": node_data.get("metadata", {}) if node_data else {},
+            }
+
+            # Get neighbors
+            neighbors = await self.get_neighbors(curr_type, curr_id, direction="outgoing", limit=20)
+            for n in neighbors:
+                edge = {
+                    "from": f"{curr_type}:{curr_id}",
+                    "to": f"{n['node_type']}:{n['node_id']}",
+                    "relation": n["relation"],
+                    "weight": n["weight"],
+                }
+                edges.append(edge)
+
+                if (n["node_type"], n["node_id"]) not in visited:
+                    queue.append((n["node_type"], n["node_id"], curr_depth + 1))
+
+        return {
+            "center": f"{node_type}:{node_id}",
+            "nodes": nodes,
+            "edges": edges,
+            "depth": depth,
+            "total_nodes": len(nodes),
+            "total_edges": len(edges),
+        }
+
+    # ── Ingestion from RAG Metadata ────────────────────────────────
+
+    async def ingest_rag_document(
+        self,
+        collection: str,
+        doc_id: str,
+        content: str,
+        metadata: dict[str, Any],
+    ) -> int:
+        """
+        Auto-extract graph relationships from a RAG document's metadata.
+        Returns number of edges created.
+        """
+        edges_created = 0
+
+        # Common metadata fields for crypto documents
+        address = metadata.get("address", "")
+        chain = metadata.get("chain", "")
+        name = metadata.get("name", "")
+        symbol = metadata.get("symbol", "")
+        labels = metadata.get("labels", [])
+        severity = metadata.get("severity", "")
+        source = metadata.get("source", "")
+
+        # ── Wallet profiles ──
+        if collection == "wallet_profiles" and address:
+            await self.add_node(
+                NodeType.WALLET,
+                address,
+                label=name or address[:10] + "...",
+                metadata={"chain": chain, "labels": labels, "source": source},
+            )
+
+            # Wallet holds tokens (if metadata has token info)
+            held_tokens = metadata.get("tokens", [])
+            for token in held_tokens if isinstance(held_tokens, list) else []:
+                token_addr = token.get("address", "") or token.get("mint", "")
+                if token_addr:
+                    await self.add_node(NodeType.TOKEN, token_addr, label=token.get("symbol", token_addr[:10]))
+                    await self.add_edge(
+                        NodeType.WALLET,
+                        address,
+                        EdgeType.HOLDS,
+                        NodeType.TOKEN,
+                        token_addr,
+                        weight=float(token.get("amount_usd", 0)) / max(float(metadata.get("balance_usd", 1)), 1) or 0.5,
+                        metadata={"symbol": token.get("symbol", ""), "chain": chain},
+                    )
+                    edges_created += 1
+
+            # Labels indicate involvement in scams/protocols
+            if isinstance(labels, list):
+                for label in labels:
+                    label_str = str(label).lower()
+                    scam_keywords = [
+                        "scam",
+                        "rug",
+                        "honeypot",
+                        "phish",
+                        "drainer",
+                        "exploit",
+                        "hack",
+                    ]
+                    if any(kw in label_str for kw in scam_keywords):
+                        await self.add_node(
+                            NodeType.SCAM,
+                            label_str,
+                            label=label_str,
+                            metadata={"type": "label", "source": source},
+                        )
+                        await self.add_edge(
+                            NodeType.WALLET,
+                            address,
+                            EdgeType.INVOLVED_IN,
+                            NodeType.SCAM,
+                            label_str,
+                            weight=0.8,
+                            metadata={"label": label_str, "chain": chain},
+                        )
+                        edges_created += 1
+
+        # ── Token analysis ──
+        elif collection == "token_analysis" and address:
+            await self.add_node(
+                NodeType.TOKEN,
+                address,
+                label=f"{name} ({symbol})" if name else address[:10],
+                metadata={"chain": chain, "symbol": symbol, "source": source},
+            )
+
+            # Deployer info
+            deployer = metadata.get("deployer", metadata.get("creator", ""))
+            if deployer:
+                await self.add_node(NodeType.WALLET, deployer, label=deployer[:10] + "...")
+                await self.add_edge(
+                    NodeType.WALLET,
+                    deployer,
+                    EdgeType.DEPLOYED,
+                    NodeType.TOKEN,
+                    address,
+                    weight=1.0,
+                    metadata={"chain": chain, "role": "deployer"},
+                )
+                edges_created += 1
+
+            # Scam flags
+            if severity in ("high", "critical"):
+                scam_label = f"{symbol or address[:8]} scam"
+                await self.add_node(
+                    NodeType.SCAM,
+                    scam_label,
+                    label=scam_label,
+                    metadata={"severity": severity, "type": "flagged_token"},
+                )
+                await self.add_edge(
+                    NodeType.TOKEN,
+                    address,
+                    EdgeType.FLAGS,
+                    NodeType.SCAM,
+                    scam_label,
+                    weight=0.9,
+                    metadata={"severity": severity},
+                )
+                edges_created += 1
+
+        # ── Known scams / scam patterns ──
+        elif collection in ("known_scams", "scam_patterns"):
+            scam_id = doc_id
+            await self.add_node(
+                NodeType.SCAM,
+                scam_id,
+                label=name or metadata.get("name", scam_id),
+                metadata={"severity": severity, "type": collection, "source": source},
+            )
+
+            # Associated addresses
+            for addr in metadata.get("addresses", []):
+                addr_str = str(addr)
+                if addr_str:
+                    await self.add_node(NodeType.ADDRESS, addr_str, label=addr_str[:10] + "...")
+                    await self.add_edge(
+                        NodeType.ADDRESS,
+                        addr_str,
+                        EdgeType.INVOLVED_IN,
+                        NodeType.SCAM,
+                        scam_id,
+                        weight=0.7,
+                        metadata={"role": "associated"},
+                    )
+                    edges_created += 1
+
+            # Associated tokens
+            for token in metadata.get("tokens", []):
+                token_addr = token.get("address", "") if isinstance(token, dict) else str(token)
+                if token_addr:
+                    await self.add_node(NodeType.TOKEN, token_addr, label=token_addr[:10])
+                    await self.add_edge(
+                        NodeType.SCAM,
+                        scam_id,
+                        EdgeType.FLAGS,
+                        NodeType.TOKEN,
+                        token_addr,
+                        weight=0.8,
+                        metadata={"severity": severity},
+                    )
+                    edges_created += 1
+
+        # ── Forensic reports ──
+        elif collection == "forensic_reports":
+            report_id = doc_id
+            await self.add_node(
+                NodeType.ENTITY,
+                report_id,
+                label=name or f"Report {report_id[:8]}",
+                metadata={"type": "forensic_report", "severity": severity},
+            )
+
+            # Connect mentioned addresses
+            for addr in metadata.get("addresses", metadata.get("wallets", [])):
+                addr_str = str(addr)
+                if addr_str:
+                    await self.add_node(NodeType.ADDRESS, addr_str, label=addr_str[:10] + "...")
+                    await self.add_edge(
+                        NodeType.ADDRESS,
+                        addr_str,
+                        EdgeType.INVOLVED_IN,
+                        NodeType.ENTITY,
+                        report_id,
+                        weight=0.6,
+                    )
+                    edges_created += 1
+
+        # ── Contract audits ──
+        elif collection == "contract_audits" and address:
+            await self.add_node(
+                NodeType.CONTRACT,
+                address,
+                label=name or address[:10],
+                metadata={"chain": chain, "source": source, "severity": severity},
+            )
+
+            # Auditor
+            auditor = metadata.get("auditor", "")
+            if auditor:
+                await self.add_node(NodeType.ENTITY, auditor, label=auditor, metadata={"type": "auditor"})
+                await self.add_edge(
+                    NodeType.ENTITY,
+                    auditor,
+                    EdgeType.ASSOCIATED,
+                    NodeType.CONTRACT,
+                    address,
+                    weight=0.5,
+                    metadata={"role": "auditor"},
+                )
+
+        return edges_created
+
+    # ── Search Enhancement ──────────────────────────────────────────
+
+    async def expand_query_entities(
+        self,
+        entities: list[dict[str, str]],
+        max_depth: int = 2,
+    ) -> list[dict[str, Any]]:
+        """
+        Given extracted entities, expand them via graph traversal.
+        Returns related entities and their relationships.
+
+        This is used by Pillar 3 (entity lookup) in three_pillar_search
+        to find not just the exact entity but its graph neighborhood.
+        """
+        expanded = []
+
+        for entity in entities:
+            entity_type = entity.get("type", "address")
+            entity_value = entity.get("value", "")
+
+            # Map entity types to graph node types
+            node_type_map = {
+                "evm_address": NodeType.WALLET,
+                "solana_address": NodeType.WALLET,
+                "token_symbol": NodeType.TOKEN,
+                "chain_name": NodeType.CHAIN,
+                "protocol_name": NodeType.ENTITY,
+            }
+            graph_type = node_type_map.get(entity_type, NodeType.ADDRESS)
+
+            # Check if this entity exists in the graph
+            node = await self.get_node(graph_type, entity_value)
+            if not node:
+                # Also try wallet/address types
+                for alt_type in [
+                    NodeType.WALLET,
+                    NodeType.TOKEN,
+                    NodeType.CONTRACT,
+                    NodeType.ADDRESS,
+                ]:
+                    node = await self.get_node(alt_type, entity_value)
+                    if node:
+                        graph_type = alt_type
+                        break
+
+            if node:
+                # Get neighbors
+                neighbors = await self.get_neighbors(
+                    graph_type,
+                    entity_value,
+                    direction="both",
+                    limit=20,
+                )
+                for n in neighbors:
+                    expanded.append(
+                        {
+                            "doc_id": f"kg:{n['node_type']}:{n['node_id']}",
+                            "id": f"kg:{n['node_type']}:{n['node_id']}",
+                            "score": n.get("weight", 0.5),
+                            "content": f"{n['label']} ({n['relation']} {n['direction']})",
+                            "pillars": ["entity"],
+                            "match_type": "knowledge_graph",
+                            "relation": n["relation"],
+                            "neighbor_type": n["node_type"],
+                            "neighbor_id": n["node_id"],
+                        }
+                    )
+
+        return expanded
+
+    # ── Statistics ──────────────────────────────────────────────────
+
+    async def stats(self) -> dict[str, Any]:
+        """Get knowledge graph statistics."""
+        r = await self._get_redis()
+
+        # Count nodes by type
+        node_counts = {}
+        for ntype in [
+            NodeType.WALLET,
+            NodeType.TOKEN,
+            NodeType.CONTRACT,
+            NodeType.SCAM,
+            NodeType.ENTITY,
+            NodeType.CHAIN,
+            NodeType.ADDRESS,
+        ]:
+            pattern = f"kg:node:{ntype}:*"
+            keys = []
+            async for key in r.scan_iter(match=pattern, count=500):
+                keys.append(key)
+            node_counts[ntype] = len(keys)
+
+        # Count edges
+        edge_keys = []
+        async for key in r.scan_iter(match="kg:edge:*", count=500):
+            edge_keys.append(key)
+
+        return {
+            "nodes": node_counts,
+            "total_nodes": sum(node_counts.values()),
+            "total_edges": len(edge_keys),
+            "node_types": list(node_counts.keys()),
+        }
+
+
+# ── Singleton ──────────────────────────────────────────────────────
+
+_kg_instance: KnowledgeGraph | None = None
+
+
+async def get_knowledge_graph() -> KnowledgeGraph:
+    """Get or create the singleton KnowledgeGraph instance."""
+    global _kg_instance
+    if _kg_instance is None:
+        _kg_instance = KnowledgeGraph()
+    return _kg_instance
+
+
+async def build_graph_from_rag(
+    collections: list[str] | None = None,
+    max_per_collection: int = 10000,
+) -> dict[str, int]:
+    """
+    One-time build: scan RAG documents and build knowledge graph edges.
+    Called on startup or as a background job.
+
+    Returns: {"edges_created": N, "nodes_created": M, "collections": {...}}
+    """
+
+    if collections is None:
+        collections = [
+            "wallet_profiles",
+            "token_analysis",
+            "known_scams",
+            "scam_patterns",
+            "forensic_reports",
+            "market_intel",
+            "contract_audits",
+            "transaction_patterns",
+        ]
+
+    kg = await get_knowledge_graph()
+    r = await kg._get_redis()  # Reuse singleton's Redis connection
+    total_edges = 0
+    total_nodes = 0
+    coll_stats = {}
+
+    for coll in collections:
+        idx_key = f"rag:idx:{coll}"
+        doc_ids = await r.smembers(idx_key)
+        if not doc_ids:
+            coll_stats[coll] = {"docs": 0, "edges": 0}
+            continue
+
+        sample = list(doc_ids)[:max_per_collection]
+        edges = 0
+
+        # Batch fetch documents
+        batch_size = 500
+        for i in range(0, len(sample), batch_size):
+            batch = sample[i : i + batch_size]
+            keys = [f"rag:{coll}:{did}" for did in batch]
+            pipe = r.pipeline()
+            for k in keys:
+                pipe.get(k)
+            raw_docs = await pipe.execute()
+
+            for did, data in zip(batch, raw_docs, strict=False):
+                if not data:
+                    continue
+                try:
+                    doc = json.loads(data)
+                    content = doc.get("content", "")
+                    metadata = doc.get("metadata", {}) or {}
+                    doc_edges = await kg.ingest_rag_document(
+                        collection=coll,
+                        doc_id=did,
+                        content=content,
+                        metadata=metadata,
+                    )
+                    edges += doc_edges
+                except (json.JSONDecodeError, Exception) as e:
+                    logger.debug(f"KG ingest failed for {coll}/{did}: {e}")
+
+        total_edges += edges
+        total_nodes += len(sample)  # approximate
+        coll_stats[coll] = {"docs": len(sample), "edges": edges}
+        logger.info(f"KG built for {coll}: {len(sample)} docs, {edges} edges")
+
+    # Connection is managed by the singleton — do not close here
+
+    return {
+        "nodes_created": total_nodes,
+        "edges_created": total_edges,
+        "collections": coll_stats,
+    }