The Knowledge Graph Stack for Agentic AI: LazyGraphRAG, Temporal Graphs, and the Architecture That Replaces Vector-Only RAG in 2026

from dataclasses import dataclass, field
from enum import Enum
import hashlib
import math
 
 
class CentralityTier(str, Enum):
    SKELETON = "skeleton"    # Top 20–30% by PageRank — full KG treatment
    BIPARTITE = "bipartite"  # Remaining 70–80% — keyword + text index only
 
 
@dataclass
class GraphNode:
    name: str
    entity_type: str
    pagerank: float = 0.0
    tier: CentralityTier = CentralityTier.BIPARTITE
    properties: dict = field(default_factory=dict)
    valid_from: str = ""    # ISO timestamp
    valid_until: str = ""   # Empty = still valid
 
    @property
    def id(self) -> str:
        raw = f"{self.entity_type}:{self.name.lower().strip()}"
        return hashlib.sha256(raw.encode()).hexdigest()[:16]
 
 
@dataclass
class TemporalEdge:
    source_id: str
    target_id: str
    relation: str
    valid_from: str          # When this fact became true
    valid_until: str = ""    # When it stopped being true (empty = current)
    confidence: float = 0.0
    source_document: str = ""
    properties: dict = field(default_factory=dict)
 
    @property
    def is_current(self) -> bool:
        return self.valid_until == ""
 
 
class SkeletonIndexer:
    """Build a KET-RAG-style skeleton graph.
    
    Only the structurally central nodes get full KG treatment.
    The rest go into a lightweight bipartite text-keyword index.
    """
 
    def __init__(self, skeleton_percentile: float = 0.3):
        self.skeleton_percentile = skeleton_percentile
        self.nodes: dict[str, GraphNode] = {}
        self.edges: list[TemporalEdge] = []
 
    def add_node(self, node: GraphNode) -> None:
        self.nodes[node.id] = node
 
    def add_edge(self, edge: TemporalEdge) -> None:
        self.edges.append(edge)
 
    def compute_skeleton(self) -> tuple[list[GraphNode], list[GraphNode]]:
        """Partition nodes into skeleton (high centrality) and bipartite (rest).
        
        Uses a simplified PageRank approximation. In production, run
        PageRank directly in Neo4j via GDS (Graph Data Science library).
        """
        if not self.nodes:
            return [], []
 
        # Approximate PageRank: count incoming edges as a proxy
        in_degree: dict[str, int] = {nid: 0 for nid in self.nodes}
        for edge in self.edges:
            if edge.target_id in in_degree and edge.is_current:
                in_degree[edge.target_id] += 1
 
        # Assign PageRank scores (simplified)
        max_degree = max(in_degree.values()) if in_degree else 1
        for nid, degree in in_degree.items():
            self.nodes[nid].pagerank = degree / max(max_degree, 1)
 
        # Partition by centrality threshold
        sorted_nodes = sorted(self.nodes.values(), key=lambda n: n.pagerank, reverse=True)
        cutoff = math.ceil(len(sorted_nodes) * self.skeleton_percentile)
 
        skeleton = sorted_nodes[:cutoff]
        bipartite = sorted_nodes[cutoff:]
 
        for node in skeleton:
            node.tier = CentralityTier.SKELETON
        for node in bipartite:
            node.tier = CentralityTier.BIPARTITE
 
        return skeleton, bipartite
 
    def build_cypher_skeleton(self, skeleton: list[GraphNode]) -> list[str]:
        """Generate Cypher MERGE statements for skeleton nodes and their edges."""
        stmts = []
        skeleton_ids = {n.id for n in skeleton}
 
        for node in skeleton:
            stmts.append(
                f"MERGE (n:{node.entity_type} {{entity_id: '{node.id}'}}) "
                f"ON CREATE SET n.name = '{node.name}', "
                f"n.pagerank = {node.pagerank}, n.tier = 'skeleton', "
                f"n.valid_from = datetime('{node.valid_from}') "
            )
 
        for edge in self.edges:
            if edge.source_id in skeleton_ids or edge.target_id in skeleton_ids:
                valid_until = f", r.valid_until = datetime('{edge.valid_until}')" if edge.valid_until else ""
                stmts.append(
                    f"MATCH (s {{entity_id: '{edge.source_id}'}}), "
                    f"(t {{entity_id: '{edge.target_id}'}}) "
                    f"MERGE (s)-[r:{edge.relation}]->(t) "
                    f"ON CREATE SET r.valid_from = datetime('{edge.valid_from}'), "
                    f"r.confidence = {edge.confidence}, "
                    f"r.source = '{edge.source_document}'"
                    f"{valid_until}"
                )
 
        return stmts

from datetime import datetime, timezone
 
 
class TemporalGraphClient:
    """Query a bi-temporal knowledge graph in Neo4j."""
 
    def __init__(self, driver):
        self._driver = driver
 
    def query_as_of(self, entity_name: str, as_of: datetime, depth: int = 2) -> list[dict]:
        """Retrieve the subgraph as it existed at a specific point in time.
        
        Only returns edges whose valid_from <= as_of and 
        (valid_until is null OR valid_until > as_of).
        """
        cypher = """
        MATCH (start {name: $entity_name})
        CALL apoc.path.subgraphAll(start, {
            maxLevel: $depth,
            relationshipFilter: null
        }) YIELD nodes, relationships
        WITH nodes, relationships
        UNWIND relationships AS r
        WHERE r.valid_from <= datetime($as_of)
          AND (r.valid_until IS NULL OR r.valid_until > datetime($as_of))
        RETURN startNode(r) AS source, type(r) AS relation, endNode(r) AS target,
               r.valid_from AS valid_from, r.valid_until AS valid_until,
               r.confidence AS confidence
        """
        with self._driver.session() as session:
            records = session.run(
                cypher,
                entity_name=entity_name,
                depth=depth,
                as_of=as_of.isoformat(),
            )
            return [dict(r) for r in records]
 
    def invalidate_edge(self, source_id: str, target_id: str, relation: str) -> None:
        """Mark an edge as no longer valid (soft delete with timestamp)."""
        cypher = """
        MATCH (s {entity_id: $source_id})-[r]->(t {entity_id: $target_id})
        WHERE type(r) = $relation AND r.valid_until IS NULL
        SET r.valid_until = datetime()
        """
        with self._driver.session() as session:
            session.run(cypher, source_id=source_id, target_id=target_id, relation=relation)
 
    def get_edge_history(self, source_id: str, target_id: str) -> list[dict]:
        """Return the full temporal history of a relationship between two entities."""
        cypher = """
        MATCH (s {entity_id: $source_id})-[r]->(t {entity_id: $target_id})
        RETURN type(r) AS relation, r.valid_from AS from, r.valid_until AS until,
               r.confidence AS confidence, r.source AS source_doc
        ORDER BY r.valid_from DESC
        """
        with self._driver.session() as session:
            records = session.run(cypher, source_id=source_id, target_id=target_id)
            return [dict(r) for r in records]

from dataclasses import dataclass
from enum import Enum
from datetime import datetime, timezone
import json
 
 
class MutationType(str, Enum):
    CREATE_NODE = "CREATE_NODE"
    UPDATE_NODE = "UPDATE_NODE"
    CREATE_EDGE = "CREATE_EDGE"
    INVALIDATE_EDGE = "INVALIDATE_EDGE"
    HEAL_EDGE = "HEAL_EDGE"          # Self-healing correction
 
 
@dataclass
class GraphEvent:
    """Immutable record of a graph mutation."""
    event_id: str
    mutation_type: MutationType
    entity_ids: list[str]
    payload: dict
    source_document: str
    agent_id: str           # Which agent/pipeline made this change
    timestamp: datetime = None
 
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now(timezone.utc)
 
    def to_json(self) -> str:
        return json.dumps({
            "event_id": self.event_id,
            "type": self.mutation_type.value,
            "entities": self.entity_ids,
            "payload": self.payload,
            "source": self.source_document,
            "agent": self.agent_id,
            "ts": self.timestamp.isoformat(),
        })
 
 
class EventStore:
    """Append-only event log for graph mutations."""
 
    def __init__(self):
        self._log: list[GraphEvent] = []
 
    def append(self, event: GraphEvent) -> None:
        self._log.append(event)
 
    def replay_since(self, since: datetime) -> list[GraphEvent]:
        """Return all events since a given timestamp."""
        return [e for e in self._log if e.timestamp >= since]
 
    def get_provenance(self, entity_id: str) -> list[GraphEvent]:
        """Full mutation history for an entity — the audit trail."""
        return [e for e in self._log if entity_id in e.entity_ids]

from dataclasses import dataclass, field
 
 
@dataclass
class DRIFTState:
    """Tracks the agent's traversal state during DRIFT search."""
    original_query: str
    global_primer: str = ""
    follow_ups: list[str] = field(default_factory=list)
    visited_entities: set = field(default_factory=set)
    accumulated_context: list[dict] = field(default_factory=list)
    depth: int = 0
    max_depth: int = 3
    sufficient: bool = False
 
 
class DRIFTSearcher:
    """DRIFT: Dynamic Reasoning and Inference with Flexible Traversal.
    
    The agent controls retrieval depth, not a fixed top-k parameter.
    """
 
    def __init__(self, graph_client, llm_client, skeleton_summaries: dict):
        self.graph = graph_client
        self.llm = llm_client
        self.summaries = skeleton_summaries  # Pre-computed community summaries
 
    async def search(self, query: str) -> DRIFTState:
        state = DRIFTState(original_query=query)
 
        # Step 1: Global primer — broad orientation from skeleton summaries
        state.global_primer = await self._generate_primer(query)
 
        # Step 2–4: Iterative deepening loop
        while not state.sufficient and state.depth < state.max_depth:
            # Generate targeted follow-up questions
            follow_ups = await self._generate_follow_ups(state)
            state.follow_ups.extend(follow_ups)
 
            # Traverse local subgraphs for each follow-up
            for fq in follow_ups:
                entities = await self._extract_entities(fq)
                for entity in entities:
                    if entity not in state.visited_entities:
                        state.visited_entities.add(entity)
                        subgraph = self.graph.query_subgraph(entity, depth=1)
                        state.accumulated_context.append({
                            "query": fq,
                            "entity": entity,
                            "subgraph": subgraph,
                            "depth": state.depth,
                        })
 
            # Evaluate sufficiency — does the agent have enough context?
            state.sufficient = await self._evaluate_sufficiency(state)
            state.depth += 1
 
        return state
 
    async def _generate_primer(self, query: str) -> str:
        """Use skeleton community summaries for broad orientation."""
        relevant_communities = self._find_relevant_communities(query)
        primer_prompt = f"""Given these community summaries from a knowledge graph:
        
{chr(10).join(relevant_communities)}
 
And this query: {query}
 
Provide a brief orientation: what topic areas and entity types are relevant?
What relationships should we explore?"""
        return await self.llm.generate(primer_prompt)
 
    async def _generate_follow_ups(self, state: DRIFTState) -> list[str]:
        """Agent generates its own follow-up questions to explore deeper."""
        prompt = f"""Original query: {state.original_query}
Global primer: {state.global_primer}
Already explored: {list(state.visited_entities)}
Current depth: {state.depth}
 
Generate 2–3 specific follow-up questions that would help answer the 
original query. Focus on relationships and entities not yet explored.
Return as a JSON array of strings."""
        response = await self.llm.generate(prompt)
        return self._parse_follow_ups(response)
 
    async def _evaluate_sufficiency(self, state: DRIFTState) -> bool:
        """Does the agent have enough context to answer the original query?"""
        prompt = f"""Original query: {state.original_query}
 
Accumulated context from {len(state.accumulated_context)} subgraph traversals
covering entities: {list(state.visited_entities)}
 
Is this context sufficient to answer the original query accurately?
Consider: Are there obvious missing relationships? Unexplored entities 
mentioned in the context? Return JSON: {{"sufficient": true/false, "reason": "..."}}"""
        response = await self.llm.generate(prompt)
        return self._parse_sufficiency(response)
 
    def _find_relevant_communities(self, query: str) -> list[str]:
        """Find skeleton community summaries relevant to the query."""
        # In production: use embedding similarity against community summary embeddings
        return list(self.summaries.values())[:5]
 
    def _parse_follow_ups(self, response: str) -> list[str]:
        import json
        try:
            return json.loads(response)
        except json.JSONDecodeError:
            return []
 
    def _parse_sufficiency(self, response: str) -> bool:
        import json
        try:
            return json.loads(response).get("sufficient", False)
        except (json.JSONDecodeError, AttributeError):
            return False
 
    async def _extract_entities(self, query: str) -> list[str]:
        """Extract entity names from a follow-up question."""
        prompt = f"Extract all named entities from this question. Return as JSON array of strings: {query}"
        response = await self.llm.generate(prompt)
        try:
            import json
            return json.loads(response)
        except json.JSONDecodeError:
            return []

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
 
 
@dataclass
class HealthCheck:
    entity_id: str
    edge_count: int
    stale_edges: int        # Edges past TTL with no reconfirmation
    conflicting_edges: int  # Multiple current edges of same type between same nodes
    orphan_score: float     # 0.0 = well-connected, 1.0 = isolated
    last_confirmed: datetime | None
 
 
class GraphHealer:
    """Detect and repair knowledge graph integrity issues."""
 
    def __init__(self, graph_client, event_store: EventStore, ttl_days: int = 90):
        self.graph = graph_client
        self.events = event_store
        self.ttl = timedelta(days=ttl_days)
 
    def audit(self) -> list[HealthCheck]:
        """Run a full health audit on the knowledge graph."""
        checks = []
        nodes = self.graph.get_all_nodes()
 
        for node in nodes:
            edges = self.graph.get_edges(node.id)
            current_edges = [e for e in edges if e.is_current]
 
            # Stale detection: edges not reconfirmed within TTL
            now = datetime.now(timezone.utc)
            stale = [
                e for e in current_edges
                if self._last_confirmation(e) < now - self.ttl
            ]
 
            # Conflict detection: multiple current edges of same type
            # between the same pair of nodes
            conflicts = self._detect_conflicts(current_edges)
 
            # Orphan detection: nodes with very low connectivity
            orphan_score = 1.0 / (1.0 + len(current_edges))
 
            last_event = self._last_event(node.id)
            checks.append(HealthCheck(
                entity_id=node.id,
                edge_count=len(current_edges),
                stale_edges=len(stale),
                conflicting_edges=len(conflicts),
                orphan_score=orphan_score,
                last_confirmed=last_event,
            ))
 
        return checks
 
    def heal_stale_edges(self, checks: list[HealthCheck]) -> int:
        """Invalidate edges that have exceeded TTL without reconfirmation."""
        healed = 0
        for check in checks:
            if check.stale_edges > 0:
                edges = self.graph.get_edges(check.entity_id)
                now = datetime.now(timezone.utc)
                for edge in edges:
                    if edge.is_current and self._last_confirmation(edge) < now - self.ttl:
                        self.graph.invalidate_edge(
                            edge.source_id, edge.target_id, edge.relation
                        )
                        self.events.append(GraphEvent(
                            event_id=f"heal-{edge.source_id}-{edge.target_id}",
                            mutation_type=MutationType.HEAL_EDGE,
                            entity_ids=[edge.source_id, edge.target_id],
                            payload={"reason": "TTL_EXPIRED", "edge": edge.relation},
                            source_document="self-healing-audit",
                            agent_id="graph-healer",
                        ))
                        healed += 1
        return healed
 
    def _detect_conflicts(self, edges: list) -> list:
        """Find edges where the same relationship type exists multiple times
        between the same pair of nodes."""
        seen = {}
        conflicts = []
        for edge in edges:
            key = (edge.source_id, edge.target_id, edge.relation)
            if key in seen:
                conflicts.append(edge)
            else:
                seen[key] = edge
        return conflicts
 
    def _last_confirmation(self, edge) -> datetime:
        """When was this edge last confirmed by a source document?"""
        events = self.events.get_provenance(edge.source_id)
        relevant = [e for e in events if edge.target_id in e.entity_ids]
        if relevant:
            return max(e.timestamp for e in relevant)
        return datetime.min.replace(tzinfo=timezone.utc)
 
    def _last_event(self, entity_id: str) -> datetime | None:
        events = self.events.get_provenance(entity_id)
        return max(e.timestamp for e in events) if events else None