klieo-memory-graph-rag 2.2.0

//! `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.
//!
//! **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).

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,
};
use std::sync::Arc;

const DEFAULT_MIN_GRAPH_HITS: usize = 1;

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

/// 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 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 => {
                    self.metrics
                        .record_graph_hit_with_candidates(candidate_ids.len() as u64);
                    return self
                        .vector
                        .recall_filtered(scope, query, k, &candidate_ids)
                        .await;
                }
                Ok(_) => { /* below threshold — fall through to vector */ }
                Err(e) => tracing::warn!(
                    error = %e,
                    "graph neighbors failed on recall; falling back to pure vector"
                ),
            }
        }
        self.metrics.record_vector_fallback();
        self.vector.recall(scope, query, k).await
    }

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