m8_graphrag.py

"""M8 — GraphRAG + HippoRAG2 Memory Layer for ARCHON.

Federated retrieval across ASI banks:
- NEXUS knowledge banks (146 GB, 25 banks)
- CYPHER security corpus (228 GB CC-BY-NC)
- AETHER memory backup
- OMOIKANE vault (PANTHEON conversations)

Architecture:
- ChromaDB (vector store, already deployed NEXUS M9)
- HippoRAG2: entity-link chunking + personalized PageRank retrieval
- Optional: Pinecone (MCP plugin already authenticated)

Ref: arxiv 2503.21322 HyperGraphRAG, GraphRAG +3.4x accuracy.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Optional

try:
    import chromadb
except ImportError:
    chromadb = None


@dataclass
class GraphRAGConfig:
    """ARCHON GraphRAG config."""

    chroma_path: str = "/workspace/archon_sft_v2/_archon_rag_db"
    collection_name: str = "archon_federated"
    embed_model: str = "BAAI/bge-m3"  # Multi-lingual SOTA
    top_k_vector: int = 20
    top_k_graph: int = 10
    pagerank_alpha: float = 0.5
    hop_max: int = 2  # Graph traversal depth
    chunk_size_tokens: int = 256


class HippoRAG2Retriever:
    """Personalized PageRank retrieval over entity-link graph.

    Phase 1 (vector): get top_k_vector seed nodes via embedding similarity.
    Phase 2 (PageRank): personalize PageRank random walk from seeds,
                        return top_k_graph by stationary distribution.
    """

    def __init__(self, cfg: GraphRAGConfig = GraphRAGConfig()):
        self.cfg = cfg
        self.client = chromadb.PersistentClient(path=cfg.chroma_path) if chromadb else None
        self.collection = None
        self.entity_graph: dict = {}  # {entity: set(linked_entities)}
        self.entity_to_chunks: dict = {}  # {entity: [chunk_ids]}

    def ingest_corpus(self, docs: list[dict]):
        """Each doc: {'id': str, 'text': str, 'metadata': dict}.

        Performs entity extraction (NER) + linkage to build the graph.
        Vector embed via bge-m3 + store ChromaDB.
        """
        if self.client is None:
            raise RuntimeError("chromadb not installed")
        if self.collection is None:
            self.collection = self.client.get_or_create_collection(self.cfg.collection_name)
        # Chunk + embed
        chunks = []
        for d in docs:
            for i, ch in enumerate(self._chunk_text(d["text"], self.cfg.chunk_size_tokens)):
                cid = f"{d['id']}::chunk{i}"
                chunks.append({"id": cid, "text": ch, "metadata": d.get("metadata", {})})
        # NER extract entities (stub — use spaCy/HF NER in prod)
        for ch in chunks:
            entities = self._extract_entities_stub(ch["text"])
            for e in entities:
                self.entity_to_chunks.setdefault(e, []).append(ch["id"])
            # Add entity co-occurrence edges
            for i, e1 in enumerate(entities):
                for e2 in entities[i + 1:]:
                    self.entity_graph.setdefault(e1, set()).add(e2)
                    self.entity_graph.setdefault(e2, set()).add(e1)
        # Embed + add to ChromaDB
        if chunks:
            self.collection.add(
                ids=[c["id"] for c in chunks],
                documents=[c["text"] for c in chunks],
                metadatas=[c["metadata"] for c in chunks],
            )

    def retrieve(self, query: str) -> list[dict]:
        """Two-stage HippoRAG2 retrieval."""
        if self.collection is None:
            return []
        # Phase 1 vector
        res = self.collection.query(query_texts=[query], n_results=self.cfg.top_k_vector)
        seed_ids = res["ids"][0]
        # Get seed entities
        seed_entities = set()
        for cid in seed_ids:
            for e, chunks in self.entity_to_chunks.items():
                if cid in chunks:
                    seed_entities.add(e)
                    break
        # Phase 2 personalized PageRank
        scores = self._pagerank(seed_entities)
        # Top entities -> chunks
        top_entities = sorted(scores.items(), key=lambda kv: -kv[1])[: self.cfg.top_k_graph]
        out = []
        seen_chunks = set()
        for ent, _ in top_entities:
            for cid in self.entity_to_chunks.get(ent, []):
                if cid not in seen_chunks:
                    seen_chunks.add(cid)
                    # Fetch text
                    ch_res = self.collection.get(ids=[cid])
                    if ch_res["documents"]:
                        out.append({
                            "id": cid,
                            "text": ch_res["documents"][0],
                            "entity": ent,
                            "score": scores[ent],
                        })
        return out[: self.cfg.top_k_graph]

    def _chunk_text(self, text: str, max_tokens: int) -> list[str]:
        # Rough word-level chunking
        words = text.split()
        return [" ".join(words[i:i + max_tokens]) for i in range(0, len(words), max_tokens)]

    def _extract_entities_stub(self, text: str) -> list[str]:
        # Stub — replace with spaCy NER or HF NER pipeline
        words = [w.strip(".,;:") for w in text.split() if w[:1].isupper() and len(w) > 3]
        return list(dict.fromkeys(words))[:8]

    def _pagerank(self, seeds: set, iters: int = 30) -> dict:
        """Personalized PageRank from seed entities."""
        scores = {e: (1.0 if e in seeds else 0.0) for e in self.entity_graph}
        alpha = self.cfg.pagerank_alpha
        for _ in range(iters):
            new = {e: (1 - alpha) * (1.0 if e in seeds else 0.0) for e in self.entity_graph}
            for e, links in self.entity_graph.items():
                if not links:
                    continue
                share = alpha * scores.get(e, 0.0) / len(links)
                for n in links:
                    new[n] = new.get(n, 0.0) + share
            scores = new
        return scores


def federate_with_nexus_banks(retriever: HippoRAG2Retriever, nexus_banks_path: str):
    """Hook to import NEXUS knowledge banks (already indexed)."""
    from pathlib import Path
    p = Path(nexus_banks_path)
    if not p.exists():
        print(f"[M8 GraphRAG] NEXUS banks not found at {nexus_banks_path}")
        return
    # Stub: load NEXUS banks_indices/*.json + bridge to ChromaDB
    print(f"[M8 GraphRAG] federation hook: {p}")


if __name__ == "__main__":
    cfg = GraphRAGConfig()
    print(f"[M8 GraphRAG] config: {cfg}")
    if chromadb is None:
        print("[M8 GraphRAG] chromadb not installed; skip live test")
    else:
        r = HippoRAG2Retriever(cfg)
        print(f"[M8 GraphRAG] retriever instantiated, ChromaDB path: {cfg.chroma_path}")