klieo_memory_graph_rag/

graph_aware.rs

//! `GraphAwareLongTerm` — drop-in `LongTermMemory` that issues graph-first
//! recall with vector fallback.
//!
//! **Authority hierarchy.** Qdrant (vector) is the source of truth for
//! every `Fact`. `KnowledgeGraph` indexing is best-effort recall infra —
//! graph failures on ingest log WARN and return the vector-minted
//! `FactId`; graph failures on recall log WARN and fall back to pure
//! vector. The composer never propagates graph errors to callers.
//!
//! **FactId provenance.** `recall` feeds the graph's `neighbors()` ids into
//! `LongTermMemory::recall_filtered`, so the graph MUST be populated only with
//! `FactId`s that exist in the vector store — i.e. via this composer's own
//! `remember` (which indexes the vector-minted id into the graph). Do NOT
//! point this at a graph an `EpisodeProjector` writes: that projector mints
//! graph-only synthetic ids that match nothing in the vector store, so its
//! facts would silently drop out of recall. Give projection its own scope.
//!
//! **Wire shape (M2 default).** Vector = `QdrantLongTerm` from
//! `klieo-memory-qdrant`; graph = `Neo4jKnowledgeGraph` from
//! `klieo-memory-graph-neo4j`; extractor = `FallbackExtractor` chaining
//! `BuiltinExtractor` (regex) → `LlmEntityExtractor` (LLM). Drop into
//! `klieo-triage` AppState behind `feature = "graph-rag"` (Task 2.4.8).
//
// Exceeds the 300-line file cap by design: single-purpose module holding
// the graph-aware long-term composer, its shared recall core, and the
// `RecallTrace` introspection type — cohesive, split would fragment one
// pipeline across files.

use crate::extractor::BuiltinExtractor;
use crate::metadata::{KEY_ENTITIES, KEY_VALID_FROM};
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};
use klieo_memory_graph::{
    EntityExtractor, EntityRef, FilterableLongTermMemory, KnowledgeGraph, RecallMetrics,
    RetrievalPath,
};
use std::sync::Arc;

const DEFAULT_MIN_GRAPH_HITS: usize = 1;

/// The observable trace behind one [`GraphAwareLongTerm::recall_traced`] call.
///
/// `ranked_facts` and `fell_back_to_vector` come from the same core that
/// [`LongTermMemory::recall`] uses and are authoritative — this trace
/// cannot diverge from production recall output. `paths` and
/// `extracted_entities` are display-only enrichment fetched via an extra
/// [`KnowledgeGraph::recall_paths`] call that production `recall` never
/// pays for; both are empty when the backend keeps the trait's default
/// no-op `recall_paths` implementation.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct RecallTrace {
    /// Entities the extractor pulled from the query text — the same input
    /// the shared core used to decide the graph-vs-vector path. Empty
    /// when the extractor failed OR legitimately found zero entities in
    /// the query.
    pub extracted_entities: Vec<EntityRef>,
    /// Per-`(entity, fact_id)` traversal hops from
    /// [`KnowledgeGraph::recall_paths`] over the extracted entities. Since
    /// `recall_paths` mirrors [`KnowledgeGraph::neighbors`] without the
    /// `min_graph_hits` fallback threshold, this can be non-empty even
    /// when `fell_back_to_vector` is `true` (a below-threshold but nonzero
    /// neighbor set). Empty when the backend keeps the trait's default
    /// no-op `recall_paths`.
    pub paths: Vec<RetrievalPath>,
    /// Authoritative ranked facts — identical to what
    /// [`LongTermMemory::recall`] returns for the same `(scope, query, k)`.
    pub ranked_facts: Vec<Fact>,
    /// `true` when this call fell through to pure vector recall (empty or
    /// failed extraction, a graph error, or a candidate count below
    /// `min_graph_hits`); `false` when the graph-filtered path was used.
    pub fell_back_to_vector: bool,
}

/// Shared result of [`GraphAwareLongTerm::recall_core`] — `facts` and
/// `fell_back_to_vector` are the two authoritative fields that both
/// `LongTermMemory::recall` and `RecallTrace` are built from.
///
/// `recall_core` records NO [`RecallMetrics`]; the caller decides
/// whether to. `graph_candidates` carries the candidate count from the
/// graph-hit arm (`0` on the fallback arm) so the metric-recording
/// caller can report it without re-running the traversal.
struct RecallCore {
    entities: Vec<EntityRef>,
    facts: Vec<Fact>,
    fell_back_to_vector: bool,
    graph_candidates: usize,
}

/// Graph-first RAG composer over a `FilterableLongTermMemory` vector
/// store + a `KnowledgeGraph` + an `EntityExtractor`.
///
/// Provenance recording is **not** the composer's responsibility —
/// wrap the `graph` argument in
/// [`crate::ProvenanceKnowledgeGraph`] at construction time to emit
/// signed chain entries on every `graph.index()` call (ADR-038).
pub struct GraphAwareLongTerm {
    vector: Arc<dyn FilterableLongTermMemory>,
    graph: Arc<dyn KnowledgeGraph>,
    extractor: Arc<dyn EntityExtractor>,
    metrics: Arc<RecallMetrics>,
    min_graph_hits: usize,
}

impl GraphAwareLongTerm {
    /// Build with the default `min_graph_hits = 1` threshold.
    pub fn new(
        vector: Arc<dyn FilterableLongTermMemory>,
        graph: Arc<dyn KnowledgeGraph>,
        extractor: Arc<dyn EntityExtractor>,
        metrics: Arc<RecallMetrics>,
    ) -> Self {
        Self {
            vector,
            graph,
            extractor,
            metrics,
            min_graph_hits: DEFAULT_MIN_GRAPH_HITS,
        }
    }

    /// Override the minimum candidate count required to take the graph
    /// path. When `graph.neighbors()` returns fewer than `n` `FactId`s,
    /// recall falls back to pure vector. Default is `1`.
    pub fn with_min_graph_hits(mut self, n: usize) -> Self {
        self.min_graph_hits = n;
        self
    }

    /// Capability-shaped builder facade — defaults the extractor to
    /// [`BuiltinExtractor`] (regex over caller hints + typed-id
    /// patterns) and the metrics to a fresh [`RecallMetrics`].
    ///
    /// ```no_run
    /// # use std::sync::Arc;
    /// # use klieo_memory_graph::{FilterableLongTermMemory, KnowledgeGraph};
    /// # use klieo_memory_graph_rag::GraphAwareLongTerm;
    /// # async fn run(vector: Arc<dyn FilterableLongTermMemory>, graph: Arc<dyn KnowledgeGraph>) {
    /// let long_term = GraphAwareLongTerm::builder().build(vector, graph);
    /// # let _ = long_term; }
    /// ```
    ///
    /// Override defaults via the builder's setters; pass the
    /// authoritative vector store + graph at [`GraphAwareLongTermBuilder::build`]
    /// time since both are required.
    pub fn builder() -> GraphAwareLongTermBuilder {
        GraphAwareLongTermBuilder::default()
    }

    /// Graph-first recall core shared by [`LongTermMemory::recall`] and
    /// [`Self::recall_traced`] — the single place that decides whether a
    /// query takes the graph-filtered path or falls back to pure vector,
    /// so the trace can never drift from what `recall` actually returns.
    ///
    /// Records no [`RecallMetrics`]: the caller owns that decision, so a
    /// debug `recall_traced` read never pollutes the production graph-hit
    /// vs vector-fallback ratio the ops endpoint reports.
    async fn recall_core(
        &self,
        scope: &Scope,
        query: &str,
        k: usize,
    ) -> Result<RecallCore, MemoryError> {
        let entities = self
            .extractor
            .extract(query, &[])
            .await
            .unwrap_or_else(|e| {
                tracing::warn!(error = %e, "extractor failed on recall; falling back to pure vector");
                Vec::new()
            });

        if !entities.is_empty() {
            match self.graph.neighbors(scope, &entities).await {
                Ok(candidate_ids) if candidate_ids.len() >= self.min_graph_hits => {
                    let facts = self
                        .vector
                        .recall_filtered(scope.clone(), query, k, &candidate_ids)
                        .await?;
                    return Ok(RecallCore {
                        entities,
                        facts,
                        fell_back_to_vector: false,
                        graph_candidates: candidate_ids.len(),
                    });
                }
                Ok(_) => { /* below threshold — fall through to vector */ }
                Err(e) => tracing::warn!(
                    error = %e,
                    "graph neighbors failed on recall; falling back to pure vector"
                ),
            }
        }
        let facts = self.vector.recall(scope.clone(), query, k).await?;
        Ok(RecallCore {
            entities,
            facts,
            fell_back_to_vector: true,
            graph_candidates: 0,
        })
    }

    /// The "why did retrieval return this?" trace behind one recall call.
    ///
    /// Runs the exact same [`Self::recall_core`] that
    /// [`LongTermMemory::recall`] uses, then additionally fetches
    /// [`KnowledgeGraph::recall_paths`] for display. Production `recall`
    /// never issues that second graph call — this method is the only path
    /// that pays for it.
    ///
    /// Records no [`RecallMetrics`] — it's a debug/introspection call, not
    /// production recall traffic.
    pub async fn recall_traced(
        &self,
        scope: Scope,
        query: &str,
        k: usize,
    ) -> Result<RecallTrace, MemoryError> {
        let core = self.recall_core(&scope, query, k).await?;
        // Display-only enrichment: a failure here must not fail the trace,
        // but log it — the sibling fallback arms in `recall_core` log theirs,
        // and a silent drop would hide backend failures from the why-lens.
        let paths = match self.graph.recall_paths(&scope, &core.entities).await {
            Ok(paths) => paths,
            Err(e) => {
                tracing::warn!(
                    ?e,
                    "recall_paths failed in recall_traced; returning empty paths"
                );
                Vec::new()
            }
        };
        Ok(RecallTrace {
            extracted_entities: core.entities,
            paths,
            ranked_facts: core.facts,
            fell_back_to_vector: core.fell_back_to_vector,
        })
    }
}

/// Builder facade for [`GraphAwareLongTerm`] — defaults the extractor
/// and metrics for the common case so callers only supply the two
/// mandatory backends (vector + graph).
#[derive(Default)]
pub struct GraphAwareLongTermBuilder {
    extractor: Option<Arc<dyn EntityExtractor>>,
    metrics: Option<Arc<RecallMetrics>>,
    min_graph_hits: Option<usize>,
}

impl GraphAwareLongTermBuilder {
    /// Override the default [`BuiltinExtractor`]. Use a
    /// [`crate::FallbackExtractor`] or [`crate::LlmEntityExtractor`] for
    /// richer extraction.
    pub fn extractor(mut self, extractor: Arc<dyn EntityExtractor>) -> Self {
        self.extractor = Some(extractor);
        self
    }

    /// Supply a shared [`RecallMetrics`] when the caller wants to read
    /// `metrics.snapshot()` from an ops/metrics endpoint. Without this,
    /// the builder allocates a fresh counter the caller cannot observe.
    pub fn metrics(mut self, metrics: Arc<RecallMetrics>) -> Self {
        self.metrics = Some(metrics);
        self
    }

    /// Override the minimum candidate count required to take the graph
    /// path. Default `1` (any graph hit wins).
    pub fn min_graph_hits(mut self, n: usize) -> Self {
        self.min_graph_hits = Some(n);
        self
    }

    /// Build the composer. `vector` is the authoritative store; `graph`
    /// is the best-effort recall index. Defaults fill in any setter
    /// that was not called.
    pub fn build(
        self,
        vector: Arc<dyn FilterableLongTermMemory>,
        graph: Arc<dyn KnowledgeGraph>,
    ) -> GraphAwareLongTerm {
        GraphAwareLongTerm {
            vector,
            graph,
            extractor: self
                .extractor
                .unwrap_or_else(|| Arc::new(BuiltinExtractor::default())),
            metrics: self
                .metrics
                .unwrap_or_else(|| Arc::new(RecallMetrics::default())),
            min_graph_hits: self.min_graph_hits.unwrap_or(DEFAULT_MIN_GRAPH_HITS),
        }
    }
}

fn extract_hints_from_metadata(metadata: &serde_json::Value) -> Vec<EntityRef> {
    metadata
        .get(KEY_ENTITIES)
        .and_then(|v| serde_json::from_value::<Vec<EntityRef>>(v.clone()).ok())
        .unwrap_or_default()
}

fn parse_valid_from(metadata: &serde_json::Value) -> Option<DateTime<Utc>> {
    metadata
        .get(KEY_VALID_FROM)
        .and_then(|v| v.as_str())
        .and_then(|s| DateTime::parse_from_rfc3339(s).ok())
        .map(|dt| dt.with_timezone(&Utc))
}

#[async_trait]
impl LongTermMemory for GraphAwareLongTerm {
    async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError> {
        // Vector is authoritative — mint the FactId first.
        let fact_id = self.vector.remember(scope.clone(), fact.clone()).await?;

        let hints = extract_hints_from_metadata(&fact.metadata);
        let valid_from = parse_valid_from(&fact.metadata);
        let entities = self
            .extractor
            .extract(&fact.text, &hints)
            .await
            .unwrap_or_else(|e| {
                tracing::warn!(error = %e, "extractor failed; indexing graph with caller hints only");
                hints.clone()
            });

        if entities.is_empty() {
            return Ok(fact_id);
        }

        if let Err(e) = self
            .graph
            .index(scope, &fact_id, &entities, &fact.text, valid_from)
            .await
        {
            tracing::warn!(
                fact_id = %fact_id,
                error = %e,
                "graph index failed; vector remains authoritative"
            );
        }
        Ok(fact_id)
    }

    async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError> {
        let core = self.recall_core(&scope, query, k).await?;
        if core.fell_back_to_vector {
            self.metrics.record_vector_fallback();
        } else {
            self.metrics
                .record_graph_hit_with_candidates(core.graph_candidates as u64);
        }
        Ok(core.facts)
    }

    async fn forget(&self, id: FactId) -> Result<(), MemoryError> {
        // Graph-side forget lands in M4 (paired with provenance chain).
        // Until then this delegates to vector only.
        self.vector.forget(id).await
    }
}