klieo_memory_graph/

traits.rs

//! Trait surface for graph-aware memory.
//!
//! Lives in `klieo-memory-graph` at 0.x (see ADR-035) — `klieo-core` stays
//! frozen at 1.x and cannot host these traits.

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::{Fact, LongTermMemory, Scope};

/// One element of a batched [`KnowledgeGraph::index_many`] call.
///
/// Mirrors the per-fact arguments of [`KnowledgeGraph::index`] so
/// backends that natively support multi-fact UNWIND / bulk-insert can
/// flush an entire batch in O(1) round-trips instead of O(N).
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct IndexEntry {
    #[allow(missing_docs)]
    pub fact_id: FactId,
    #[allow(missing_docs)]
    pub entities: Vec<EntityRef>,
    /// Fact text — passed for parity with single `index`. Backends
    /// that don't store the text on the graph node can ignore it.
    pub text: String,
    /// World-time validity. `None` defaults to `recorded_at`.
    pub valid_from: Option<DateTime<Utc>>,
}

impl IndexEntry {
    #[allow(missing_docs)]
    pub fn new(
        fact_id: FactId,
        entities: Vec<EntityRef>,
        text: impl Into<String>,
        valid_from: Option<DateTime<Utc>>,
    ) -> Self {
        Self {
            fact_id,
            entities,
            text: text.into(),
            valid_from,
        }
    }
}

/// Structured entity graph over stored facts.
///
/// Implementations index facts as `MENTIONED_IN` edges from `Entity` nodes
/// to `FactRef` nodes, plus `CO_OCCURS` edges between entities seen in
/// the same fact. `neighbors()` returns fact ids reachable via direct
/// `MENTIONED_IN` plus 1-hop `CO_OCCURS` sibling traversal — see
/// [`Self::neighbors`] for the precise contract.
#[async_trait]
pub trait KnowledgeGraph: Send + Sync {
    /// Index a stored fact's entities into the graph.
    ///
    /// `valid_from` records when the knowledge was valid in the world.
    /// `None` defaults to `recorded_at` (now). Pass `known_epoch()` for
    /// all-time knowledge (regulatory texts, baseline facts).
    ///
    /// ### Upsert semantics
    ///
    /// `index()` is **additive**, not replacing — re-indexing the same
    /// `FactId` with a different entity set adds new `MENTIONED_IN` edges
    /// and increments `CO_OCCURS` counts for any newly-co-occurring pairs.
    /// **Existing edges are not removed.** Callers that need replace
    /// semantics must call `forget(fact_id)` first (M4). M1 does not ship
    /// `forget()`; the entire fact-set is append-only until M4.
    async fn index(
        &self,
        scope: Scope,
        fact_id: &FactId,
        entities: &[EntityRef],
        text: &str,
        valid_from: Option<DateTime<Utc>>,
    ) -> Result<(), MemoryError>;

    /// Index a batch of facts under the same [`Scope`] in one call.
    ///
    /// Same upsert semantics as [`Self::index`] applied per entry.
    /// Default impl loops on [`Self::index`] sequentially so
    /// existing backends remain correct without source change;
    /// backends with native bulk-insert (e.g. Neo4j UNWIND) override
    /// for O(1) round-trip per batch.
    ///
    /// ### Skip-empty contract
    ///
    /// Entries whose `entities` is empty are skipped — no graph
    /// write, no observable side effect. Wrappers (e.g.
    /// `ProvenanceKnowledgeGraph`) MUST filter the same way before
    /// emitting side-effect ledgers so audit state aligns with
    /// graph state.
    ///
    /// ### Failure
    ///
    /// Fails fast on the first entry that errors — partial progress
    /// is left committed because [`Self::index`] is additive (no
    /// rollback across entries). Callers that need atomic batches
    /// must wrap the call in a backend-specific transaction.
    async fn index_many(&self, scope: Scope, batch: &[IndexEntry]) -> Result<(), MemoryError> {
        for entry in batch {
            if entry.entities.is_empty() {
                continue;
            }
            self.index(
                scope.clone(),
                &entry.fact_id,
                &entry.entities,
                &entry.text,
                entry.valid_from,
            )
            .await?;
        }
        Ok(())
    }

    /// Return `FactId`s of facts reachable from `entities` via:
    /// (a) direct `MENTIONED_IN` edges, and
    /// (b) 1-hop `CO_OCCURS` siblings' `MENTIONED_IN` edges.
    ///
    /// Depth is fixed at 1-hop CO_OCCURS for M1; a `hops`/`max_hops`
    /// parameter is deliberately omitted until M5 introduces PathRAG
    /// traversal depth control (no impl currently honors it, so the
    /// signature stays honest about what callers get).
    async fn neighbors(
        &self,
        scope: &Scope,
        entities: &[EntityRef],
    ) -> Result<Vec<FactId>, MemoryError>;

    /// Return one [`crate::RetrievalPath`] per `(entity, fact_id)`
    /// edge reachable from `entities` via the same traversal as
    /// [`Self::neighbors`]. Each hop carries `chain_entry: None`;
    /// the `ProvenanceKnowledgeGraph` wrapper attaches provenance
    /// after the inner graph returns.
    ///
    /// Default impl returns `Ok(Vec::new())` so existing backends
    /// continue to compile against the 0.2 trait surface without
    /// code change until they opt in to per-path output.
    async fn recall_paths(
        &self,
        scope: &Scope,
        entities: &[EntityRef],
    ) -> Result<Vec<crate::RetrievalPath>, MemoryError> {
        let _ = (scope, entities);
        Ok(Vec::new())
    }

    /// Enumerate up to `limit` nodes of `scope`'s graph for browsing.
    ///
    /// Read-only, retrieval-agnostic — unlike [`Self::neighbors`] it needs
    /// no seed entities. `GraphView::truncated` is set when the scope held
    /// more nodes than `limit`. When truncated, the retained subset is
    /// backend-defined but stable for a given graph state (the in-memory
    /// backend keeps lowest node-index first). Default impl returns an empty
    /// view so backends compile unchanged until they opt in.
    async fn subgraph(&self, scope: &Scope, limit: usize) -> Result<crate::GraphView, MemoryError> {
        let _ = (scope, limit);
        Ok(crate::GraphView::default())
    }

    /// How many facts mention each entity under `scope`, keyed by entity
    /// name — the corpus-wide document-frequency tally.
    ///
    /// Answers "is this term generic in this corpus", which every consumer
    /// that wants to tell a specific term from a corpus-generic one needs.
    /// Without it the only route is walking [`Self::subgraph`] and tallying
    /// client-side, and `subgraph` takes a `limit` that returns a truncation
    /// error rather than data — so the client has to pick a cap and hope.
    /// A backend can count without materialising the view, which is the part
    /// the client-side version cannot do.
    ///
    /// Names are the normalised (lowercase) [`EntityRef::name`], so an entity
    /// indexed under several types is counted once per mention across all of
    /// them. Entities with no mentions are absent, not zero.
    ///
    /// The default fails loudly rather than guessing: a silent `Ok(empty)`
    /// reads as "nothing is generic in this corpus" and would let every
    /// consumer's threshold pass on a backend that never counted anything.
    async fn all_entity_frequencies(
        &self,
        scope: &Scope,
    ) -> Result<std::collections::HashMap<String, usize>, MemoryError> {
        let _ = scope;
        Err(MemoryError::Store(
            "backend does not implement all_entity_frequencies".into(),
        ))
    }

    /// [`Self::all_entity_frequencies`] restricted to `names`.
    ///
    /// Defaults to filtering the full tally so a backend gets both from one
    /// implementation; override when the store can count the named entities
    /// directly, which is cheaper for a short list over a large corpus.
    /// A name with no mentions maps to `0` rather than being absent — the
    /// caller asked about it, so it gets an answer.
    async fn entity_frequencies(
        &self,
        scope: &Scope,
        names: &[&str],
    ) -> Result<std::collections::HashMap<String, usize>, MemoryError> {
        let all = self.all_entity_frequencies(scope).await?;
        Ok(names
            .iter()
            .map(|name| ((*name).to_owned(), all.get(*name).copied().unwrap_or(0)))
            .collect())
    }

    /// Remove a fact's graph-side index entries. Paired with a
    /// `Custom("MemoryForget")` provenance event by
    /// `ProvenanceKnowledgeGraph` so the audit chain records the
    /// deletion even though the chain itself stays append-only.
    ///
    /// Default impl is a no-op; backends that support deletion
    /// override. Calling on a fact that was never indexed must not
    /// be an error.
    async fn forget(&self, scope: &Scope, fact_id: &FactId) -> Result<(), MemoryError> {
        let _ = (scope, fact_id);
        Ok(())
    }
}

/// `LongTermMemory` extension that supports candidate-set filtering.
///
/// Used by `GraphAwareLongTerm` (M2) to restrict vector recall to the
/// fact ids surfaced by the graph traversal.
#[async_trait]
pub trait FilterableLongTermMemory: LongTermMemory {
    /// Top-`k` semantic recall restricted to the given `candidate_ids`.
    ///
    /// Returns an empty `Vec` when `candidate_ids` is empty — callers
    /// should treat empty candidates as "fall back to pure vector".
    async fn recall_filtered(
        &self,
        scope: Scope,
        query: &str,
        k: usize,
        candidate_ids: &[FactId],
    ) -> Result<Vec<Fact>, MemoryError>;

    /// Identifier of the embedder this store was initialised with.
    /// Used by `recall_filtered_checked()` (M2 concrete impls) to
    /// hard-fail on cross-embedder queries — re-indexing required
    /// before mixing embedder versions.
    fn embedder_id(&self) -> &str;
}

/// Extracts typed entity references from text.
///
/// Implementations include `BuiltinExtractor` (regex, M2),
/// `LlmEntityExtractor` (M2), `FallbackExtractor` (chain primary→secondary
/// on empty, M2). Tests use a `FakeExtractor` returning a fixed set.
#[async_trait]
pub trait EntityExtractor: Send + Sync {
    /// Extract entities from `text`, merging with caller-supplied `hints`.
    ///
    /// Hints take precedence — extractor implementations dedupe by
    /// `(EntityType::as_str(), name)` and never drop a hint silently.
    async fn extract(&self, text: &str, hints: &[EntityRef])
        -> Result<Vec<EntityRef>, MemoryError>;

    /// Stable identity of *this* extractor, including anything it wraps.
    ///
    /// Exists for memoising decorators. A cache keyed on the text, the hints
    /// and a caller-supplied model string but **not** on the extractor it
    /// decorates serves the previous extractor's entries when the extractor
    /// is swapped: the second arm of an A/B reads the first arm's cache, both
    /// arms report identical numbers at full speed, and the run reads as
    /// "the extractor makes no difference" rather than "the second arm never
    /// ran". Fold this into the key instead.
    ///
    /// The default is the concrete type name, which distinguishes swaps
    /// between extractor *types*. Override when two instances of the same
    /// type differ in a way that changes output — a different model, a
    /// different prompt, a different regex table — and compose the inner
    /// ids in decorators so the whole chain is named.
    fn id(&self) -> std::borrow::Cow<'static, str> {
        std::borrow::Cow::Borrowed(std::any::type_name::<Self>())
    }
}