Skip to main content

klieo_memory_graph/
traits.rs

1//! Trait surface for graph-aware memory.
2//!
3//! Lives in `klieo-memory-graph` at 0.x (see ADR-035) — `klieo-core` stays
4//! frozen at 1.x and cannot host these traits.
5
6use crate::types::EntityRef;
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use klieo_core::error::MemoryError;
10use klieo_core::ids::FactId;
11use klieo_core::memory::{Fact, LongTermMemory, Scope};
12
13/// One element of a batched [`KnowledgeGraph::index_many`] call.
14///
15/// Mirrors the per-fact arguments of [`KnowledgeGraph::index`] so
16/// backends that natively support multi-fact UNWIND / bulk-insert can
17/// flush an entire batch in O(1) round-trips instead of O(N).
18#[derive(Debug, Clone)]
19#[non_exhaustive]
20pub struct IndexEntry {
21    #[allow(missing_docs)]
22    pub fact_id: FactId,
23    #[allow(missing_docs)]
24    pub entities: Vec<EntityRef>,
25    /// Fact text — passed for parity with single `index`. Backends
26    /// that don't store the text on the graph node can ignore it.
27    pub text: String,
28    /// World-time validity. `None` defaults to `recorded_at`.
29    pub valid_from: Option<DateTime<Utc>>,
30}
31
32impl IndexEntry {
33    #[allow(missing_docs)]
34    pub fn new(
35        fact_id: FactId,
36        entities: Vec<EntityRef>,
37        text: impl Into<String>,
38        valid_from: Option<DateTime<Utc>>,
39    ) -> Self {
40        Self {
41            fact_id,
42            entities,
43            text: text.into(),
44            valid_from,
45        }
46    }
47}
48
49/// Structured entity graph over stored facts.
50///
51/// Implementations index facts as `MENTIONED_IN` edges from `Entity` nodes
52/// to `FactRef` nodes, plus `CO_OCCURS` edges between entities seen in
53/// the same fact. `neighbors()` returns fact ids reachable via direct
54/// `MENTIONED_IN` plus 1-hop `CO_OCCURS` sibling traversal — see
55/// [`Self::neighbors`] for the precise contract.
56#[async_trait]
57pub trait KnowledgeGraph: Send + Sync {
58    /// Index a stored fact's entities into the graph.
59    ///
60    /// `valid_from` records when the knowledge was valid in the world.
61    /// `None` defaults to `recorded_at` (now). Pass `known_epoch()` for
62    /// all-time knowledge (regulatory texts, baseline facts).
63    ///
64    /// ### Upsert semantics
65    ///
66    /// `index()` is **additive**, not replacing — re-indexing the same
67    /// `FactId` with a different entity set adds new `MENTIONED_IN` edges
68    /// and increments `CO_OCCURS` counts for any newly-co-occurring pairs.
69    /// **Existing edges are not removed.** Callers that need replace
70    /// semantics must call `forget(fact_id)` first (M4). M1 does not ship
71    /// `forget()`; the entire fact-set is append-only until M4.
72    async fn index(
73        &self,
74        scope: Scope,
75        fact_id: &FactId,
76        entities: &[EntityRef],
77        text: &str,
78        valid_from: Option<DateTime<Utc>>,
79    ) -> Result<(), MemoryError>;
80
81    /// Index a batch of facts under the same [`Scope`] in one call.
82    ///
83    /// Same upsert semantics as [`Self::index`] applied per entry.
84    /// Default impl loops on [`Self::index`] sequentially so
85    /// existing backends remain correct without source change;
86    /// backends with native bulk-insert (e.g. Neo4j UNWIND) override
87    /// for O(1) round-trip per batch.
88    ///
89    /// ### Skip-empty contract
90    ///
91    /// Entries whose `entities` is empty are skipped — no graph
92    /// write, no observable side effect. Wrappers (e.g.
93    /// `ProvenanceKnowledgeGraph`) MUST filter the same way before
94    /// emitting side-effect ledgers so audit state aligns with
95    /// graph state.
96    ///
97    /// ### Failure
98    ///
99    /// Fails fast on the first entry that errors — partial progress
100    /// is left committed because [`Self::index`] is additive (no
101    /// rollback across entries). Callers that need atomic batches
102    /// must wrap the call in a backend-specific transaction.
103    async fn index_many(&self, scope: Scope, batch: &[IndexEntry]) -> Result<(), MemoryError> {
104        for entry in batch {
105            if entry.entities.is_empty() {
106                continue;
107            }
108            self.index(
109                scope.clone(),
110                &entry.fact_id,
111                &entry.entities,
112                &entry.text,
113                entry.valid_from,
114            )
115            .await?;
116        }
117        Ok(())
118    }
119
120    /// Return `FactId`s of facts reachable from `entities` via:
121    /// (a) direct `MENTIONED_IN` edges, and
122    /// (b) 1-hop `CO_OCCURS` siblings' `MENTIONED_IN` edges.
123    ///
124    /// Depth is fixed at 1-hop CO_OCCURS for M1; a `hops`/`max_hops`
125    /// parameter is deliberately omitted until M5 introduces PathRAG
126    /// traversal depth control (no impl currently honors it, so the
127    /// signature stays honest about what callers get).
128    async fn neighbors(
129        &self,
130        scope: &Scope,
131        entities: &[EntityRef],
132    ) -> Result<Vec<FactId>, MemoryError>;
133
134    /// Return one [`crate::RetrievalPath`] per `(entity, fact_id)`
135    /// edge reachable from `entities` via the same traversal as
136    /// [`Self::neighbors`]. Each hop carries `chain_entry: None`;
137    /// the `ProvenanceKnowledgeGraph` wrapper attaches provenance
138    /// after the inner graph returns.
139    ///
140    /// Default impl returns `Ok(Vec::new())` so existing backends
141    /// continue to compile against the 0.2 trait surface without
142    /// code change until they opt in to per-path output.
143    async fn recall_paths(
144        &self,
145        scope: &Scope,
146        entities: &[EntityRef],
147    ) -> Result<Vec<crate::RetrievalPath>, MemoryError> {
148        let _ = (scope, entities);
149        Ok(Vec::new())
150    }
151
152    /// Enumerate up to `limit` nodes of `scope`'s graph for browsing.
153    ///
154    /// Read-only, retrieval-agnostic — unlike [`Self::neighbors`] it needs
155    /// no seed entities. `GraphView::truncated` is set when the scope held
156    /// more nodes than `limit`. When truncated, the retained subset is
157    /// backend-defined but stable for a given graph state (the in-memory
158    /// backend keeps lowest node-index first). Default impl returns an empty
159    /// view so backends compile unchanged until they opt in.
160    async fn subgraph(&self, scope: &Scope, limit: usize) -> Result<crate::GraphView, MemoryError> {
161        let _ = (scope, limit);
162        Ok(crate::GraphView::default())
163    }
164
165    /// Remove a fact's graph-side index entries. Paired with a
166    /// `Custom("MemoryForget")` provenance event by
167    /// `ProvenanceKnowledgeGraph` so the audit chain records the
168    /// deletion even though the chain itself stays append-only.
169    ///
170    /// Default impl is a no-op; backends that support deletion
171    /// override. Calling on a fact that was never indexed must not
172    /// be an error.
173    async fn forget(&self, scope: &Scope, fact_id: &FactId) -> Result<(), MemoryError> {
174        let _ = (scope, fact_id);
175        Ok(())
176    }
177}
178
179/// `LongTermMemory` extension that supports candidate-set filtering.
180///
181/// Used by `GraphAwareLongTerm` (M2) to restrict vector recall to the
182/// fact ids surfaced by the graph traversal.
183#[async_trait]
184pub trait FilterableLongTermMemory: LongTermMemory {
185    /// Top-`k` semantic recall restricted to the given `candidate_ids`.
186    ///
187    /// Returns an empty `Vec` when `candidate_ids` is empty — callers
188    /// should treat empty candidates as "fall back to pure vector".
189    async fn recall_filtered(
190        &self,
191        scope: Scope,
192        query: &str,
193        k: usize,
194        candidate_ids: &[FactId],
195    ) -> Result<Vec<Fact>, MemoryError>;
196
197    /// Identifier of the embedder this store was initialised with.
198    /// Used by `recall_filtered_checked()` (M2 concrete impls) to
199    /// hard-fail on cross-embedder queries — re-indexing required
200    /// before mixing embedder versions.
201    fn embedder_id(&self) -> &str;
202}
203
204/// Extracts typed entity references from text.
205///
206/// Implementations include `BuiltinExtractor` (regex, M2),
207/// `LlmEntityExtractor` (M2), `FallbackExtractor` (chain primary→secondary
208/// on empty, M2). Tests use a `FakeExtractor` returning a fixed set.
209#[async_trait]
210pub trait EntityExtractor: Send + Sync {
211    /// Extract entities from `text`, merging with caller-supplied `hints`.
212    ///
213    /// Hints take precedence — extractor implementations dedupe by
214    /// `(EntityType::as_str(), name)` and never drop a hint silently.
215    async fn extract(&self, text: &str, hints: &[EntityRef])
216        -> Result<Vec<EntityRef>, MemoryError>;
217}