klieo_memory_graph/

in_memory.rs

//! `InMemoryGraph` — petgraph `StableGraph`-backed [`KnowledgeGraph`] impl.
//!
//! For tests and the hello-agent M1 spike only — not for production use.
//! Production deployments wire `klieo-memory-graph-neo4j::Neo4jKnowledgeGraph`
//! (M2). `StableGraph` is used over `Graph` because handle stability survives
//! removal — needed once `forget()` lands.

use crate::path::{PathHop, RetrievalPath};
use crate::traits::KnowledgeGraph;
use crate::types::EntityRef;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_core::error::MemoryError;
use klieo_core::ids::FactId;
use klieo_core::memory::Scope;
use petgraph::stable_graph::{NodeIndex, StableGraph};
use petgraph::Undirected;
use std::collections::{HashMap, HashSet};
use std::sync::Mutex;

/// `Node` discriminator carrying only the identifier needed to surface the
/// neighbor's identity; the scope is enforced via the index keys, not the
/// payload. The `FactId` payload is the actual return value of `neighbors()`,
/// so the Entity variant intentionally carries no payload at this milestone.
#[derive(Debug, Clone)]
enum Node {
    Entity,
    FactRef(FactId),
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ScopeKey {
    kind: &'static str,
    value: String,
}

impl From<&Scope> for ScopeKey {
    fn from(scope: &Scope) -> Self {
        match scope {
            Scope::Workspace(s) => Self {
                kind: "workspace",
                value: s.clone(),
            },
            Scope::Agent(s) => Self {
                kind: "agent",
                value: s.clone(),
            },
            Scope::Global => Self {
                kind: "global",
                value: String::new(),
            },
        }
    }
}

#[derive(Debug, Clone)]
enum Edge {
    MentionedIn,
    CoOccurs { count: u32 },
}

struct Inner {
    graph: StableGraph<Node, Edge, Undirected>,
    entity_idx: HashMap<(String, String, ScopeKey), NodeIndex>,
    /// Scope-keyed FactRef index. Two facts with the same `FactId` in
    /// different scopes get distinct nodes — closing the cross-scope leak
    /// the prior un-scoped key allowed via shared-FactRef neighbors.
    factref_idx: HashMap<(String, ScopeKey), NodeIndex>,
}

/// In-memory [`KnowledgeGraph`] backed by a petgraph `StableGraph`.
///
/// `StableGraph` preserves node/edge indices after removal — required for
/// future `forget()` support without re-indexing every entry. The single
/// `Mutex<Inner>` is acceptable for the spike scale; production traffic
/// goes to `Neo4jKnowledgeGraph` (M2).
pub struct InMemoryGraph {
    inner: Mutex<Inner>,
}

impl Default for InMemoryGraph {
    fn default() -> Self {
        Self {
            inner: Mutex::new(Inner {
                graph: StableGraph::default(),
                entity_idx: HashMap::new(),
                factref_idx: HashMap::new(),
            }),
        }
    }
}

#[async_trait]
impl KnowledgeGraph for InMemoryGraph {
    async fn index(
        &self,
        scope: Scope,
        fact_id: &FactId,
        entities: &[EntityRef],
        _text: &str,
        _valid_from: Option<DateTime<Utc>>,
    ) -> Result<(), MemoryError> {
        if entities.is_empty() {
            tracing::debug!(%fact_id, "index called with empty entities; no-op");
            return Ok(());
        }
        let sk = ScopeKey::from(&scope);
        // No `.await` while the std Mutex guard is held — sound for Send across the async boundary.
        let mut guard = self.inner.lock().map_err(|_| {
            tracing::error!(operation = "InMemoryGraph::index", "mutex poisoned");
            MemoryError::Store("InMemoryGraph mutex poisoned".into())
        })?;
        // Split-borrow: destructure once so closures don't reborrow the guard.
        let Inner {
            graph,
            entity_idx,
            factref_idx,
        } = &mut *guard;

        let fact_node = *factref_idx
            .entry((fact_id.to_string(), sk.clone()))
            .or_insert_with(|| graph.add_node(Node::FactRef(fact_id.clone())));

        let mut entity_nodes: Vec<NodeIndex> = Vec::with_capacity(entities.len());
        for entity in entities {
            let key = (
                entity.entity_type.as_str().to_owned(),
                entity.name.clone(),
                sk.clone(),
            );
            let ent_node = *entity_idx
                .entry(key)
                .or_insert_with(|| graph.add_node(Node::Entity));
            entity_nodes.push(ent_node);
            if !graph.contains_edge(ent_node, fact_node) {
                graph.add_edge(ent_node, fact_node, Edge::MentionedIn);
            }
        }

        for i in 0..entity_nodes.len() {
            for j in (i + 1)..entity_nodes.len() {
                let (a, b) = (entity_nodes[i], entity_nodes[j]);
                if let Some(edge) = graph.find_edge(a, b) {
                    if let Some(Edge::CoOccurs { count }) = graph.edge_weight_mut(edge) {
                        *count += 1;
                    }
                } else {
                    graph.add_edge(a, b, Edge::CoOccurs { count: 1 });
                }
            }
        }
        Ok(())
    }

    async fn neighbors(
        &self,
        scope: &Scope,
        entities: &[EntityRef],
    ) -> Result<Vec<FactId>, MemoryError> {
        if entities.is_empty() {
            tracing::debug!("neighbors called with empty entities; no-op");
            return Ok(Vec::new());
        }
        let sk = ScopeKey::from(scope);
        let guard = self.inner.lock().map_err(|_| {
            tracing::error!(operation = "InMemoryGraph::neighbors", "mutex poisoned");
            MemoryError::Store("InMemoryGraph mutex poisoned".into())
        })?;
        let Inner {
            graph, entity_idx, ..
        } = &*guard;
        let mut seen: HashSet<FactId> = HashSet::new();

        for entity in entities {
            let key = (
                entity.entity_type.as_str().to_owned(),
                entity.name.clone(),
                sk.clone(),
            );
            let Some(&ent_node) = entity_idx.get(&key) else {
                continue;
            };
            collect_reachable_fact_ids(graph, ent_node, &mut seen);
        }

        Ok(seen.into_iter().collect())
    }

    async fn recall_paths(
        &self,
        scope: &Scope,
        entities: &[EntityRef],
    ) -> Result<Vec<RetrievalPath>, MemoryError> {
        if entities.is_empty() {
            return Ok(Vec::new());
        }
        let sk = ScopeKey::from(scope);
        let guard = self.inner.lock().map_err(|_| {
            tracing::error!(operation = "InMemoryGraph::recall_paths", "mutex poisoned");
            MemoryError::Store("InMemoryGraph mutex poisoned".into())
        })?;
        let Inner {
            graph, entity_idx, ..
        } = &*guard;

        let mut paths: Vec<RetrievalPath> = Vec::new();
        let mut seen: HashSet<(String, String, FactId)> = HashSet::new();

        for entity in entities {
            let key = (
                entity.entity_type.as_str().to_owned(),
                entity.name.clone(),
                sk.clone(),
            );
            let Some(&ent_node) = entity_idx.get(&key) else {
                continue;
            };
            let mut fact_ids: HashSet<FactId> = HashSet::new();
            collect_reachable_fact_ids(graph, ent_node, &mut fact_ids);
            for fid in fact_ids {
                let dedupe_key = (
                    entity.entity_type.as_str().to_owned(),
                    entity.name.clone(),
                    fid.clone(),
                );
                if !seen.insert(dedupe_key) {
                    continue;
                }
                paths.push(RetrievalPath {
                    hops: vec![PathHop::new(entity.clone(), fid)],
                });
            }
        }
        Ok(paths)
    }

    async fn forget(&self, scope: &Scope, fact_id: &FactId) -> Result<(), MemoryError> {
        let sk = ScopeKey::from(scope);
        let mut guard = self.inner.lock().map_err(|_| {
            tracing::error!(operation = "InMemoryGraph::forget", "mutex poisoned");
            MemoryError::Store("InMemoryGraph mutex poisoned".into())
        })?;
        let Inner {
            graph, factref_idx, ..
        } = &mut *guard;

        let Some(fact_node) = factref_idx.remove(&(fact_id.to_string(), sk)) else {
            return Ok(());
        };

        // CO_OCCURS edges live entity↔entity and survive — co-occurrence
        // is a historical fact about the original index call.
        let fact_edges: Vec<_> = graph
            .edges(fact_node)
            .map(|e| petgraph::visit::EdgeRef::id(&e))
            .collect();
        for edge in fact_edges {
            graph.remove_edge(edge);
        }
        graph.remove_node(fact_node);
        Ok(())
    }
}

/// Equivalent to the Neo4j Cypher `UNION` in `klieo-memory-graph-neo4j`,
/// so both backends return the same fact-id set for any given entry entity.
fn collect_reachable_fact_ids(
    graph: &StableGraph<Node, Edge, Undirected>,
    entry_entity: NodeIndex,
    seen: &mut HashSet<FactId>,
) {
    for neighbor in graph.neighbors(entry_entity) {
        match &graph[neighbor] {
            Node::FactRef(fid) => {
                seen.insert(fid.clone());
            }
            Node::Entity => {
                for deeper in graph.neighbors(neighbor) {
                    if let Node::FactRef(fid) = &graph[deeper] {
                        seen.insert(fid.clone());
                    }
                }
            }
        }
    }
}