klieo-memory-graph 3.5.0

//! 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())
    }

    /// 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>;
}