klieo 3.3.1

//! Configuration for the one-line `App::graph_rag` facade.
//!
//! Two tiers:
//!
//! - [`GraphRagConfig::in_memory`] — zero-infrastructure demo/dev tier:
//!   in-process graph + vector, no external services.
//! - [`GraphRagConfig::remote`] — production tier: Qdrant vector store +
//!   Neo4j knowledge graph.
//!
//! All capabilities default ON; chainable opt-out methods disable
//! individual components.

use std::sync::Arc;

use klieo_core::error::Error;
use klieo_core::llm::LlmClient;
use klieo_core::memory::{LongTermMemory, Scope};
use klieo_embed_common::Embedder;
use klieo_memory_graph::{
    EntityExtractor, FilterableLongTermMemory, KnowledgeGraph, RecallMetrics,
};
use klieo_memory_graph_rag::{
    BuiltinExtractor, FallbackExtractor, GraphAwareLongTerm, InMemoryFilterableLongTerm,
    LlmEntityExtractor, ProvenanceKnowledgeGraph,
};

/// Neo4j connection config for [`GraphRagConfig::remote`], re-exported so
/// callers configure the remote backend through `klieo` without taking a
/// direct dependency on the `klieo-memory-neo4j` adapter crate.
pub use klieo_memory_neo4j::Neo4jConfig;

const DEFAULT_MIN_GRAPH_HITS: usize = 1;
const RESERVED_PROJECTION_SCOPE: &str = "__klieo_graph_projection";

/// Provenance actor recorded on every emitted event when the App's
/// graph-aware long-term port writes through the provenance wrapper.
const PROVENANCE_ACTOR: &str = "klieo-app-graph-rag";

/// Default on-disk provenance database used for the remote backend when
/// no explicit sqlite path is supplied by the App builder.
const DEFAULT_PROVENANCE_DB_PATH: &str = "klieo-provenance.sqlite";

/// Default fastembed model identifier used when no custom embedder is provided.
pub(crate) const FASTEMBED_ID: &str = "fastembed:all-minilm-l6-v2";

/// How entities are extracted from text before graph indexing.
#[non_exhaustive]
pub enum ExtractorMode {
    /// Run the built-in regex/NLP extractor first, then call the active LLM
    /// to fill gaps the built-in missed.
    BuiltinThenLlm,
    /// Run the built-in extractor only; skip the LLM fallback pass.
    BuiltinOnly,
    /// Delegate entirely to a caller-supplied extractor implementation.
    Custom(Arc<dyn EntityExtractor>),
}

pub(crate) enum Backend {
    InMemory,
    Remote {
        qdrant_url: String,
        neo4j: klieo_memory_neo4j::Neo4jConfig,
    },
}

/// Configuration for `App::graph_rag_with`.
///
/// Build via [`in_memory`][Self::in_memory] (zero-infra) or
/// [`remote`][Self::remote] (production backends), then opt out of
/// individual capabilities with the chainable knob methods.
#[non_exhaustive]
pub struct GraphRagConfig {
    pub(crate) backend: Backend,
    pub(crate) embedder: Option<(Arc<dyn Embedder>, String)>,
    pub(crate) provenance: bool,
    pub(crate) projector: bool,
    pub(crate) extractor_mode: ExtractorMode,
    pub(crate) projection_scope: Scope,
    pub(crate) min_graph_hits: usize,
    pub(crate) provenance_actor: String,
}

impl GraphRagConfig {
    fn defaults(backend: Backend) -> Self {
        Self {
            backend,
            embedder: None,
            provenance: true,
            projector: true,
            extractor_mode: ExtractorMode::BuiltinThenLlm,
            projection_scope: Scope::Workspace(RESERVED_PROJECTION_SCOPE.to_string()),
            min_graph_hits: DEFAULT_MIN_GRAPH_HITS,
            provenance_actor: PROVENANCE_ACTOR.to_string(),
        }
    }

    /// Zero-infrastructure tier: in-process graph + vector, no external
    /// services required. Suitable for demos and local development.
    pub fn in_memory() -> Self {
        Self::defaults(Backend::InMemory)
    }

    /// Production tier: Qdrant at `qdrant_url` for vector recall, Neo4j via
    /// `neo4j` for the knowledge graph. The remote backend writes provenance to
    /// an on-disk sqlite database; if the App provides no sqlite path it
    /// defaults to `klieo-provenance.sqlite` in the process working directory.
    pub fn remote(qdrant_url: impl Into<String>, neo4j: klieo_memory_neo4j::Neo4jConfig) -> Self {
        Self::defaults(Backend::Remote {
            qdrant_url: qdrant_url.into(),
            neo4j,
        })
    }

    /// Override the default fastembed embedder with `embedder`, identified by
    /// `id` for model-registry lookup and cost attribution.
    pub fn embedder(mut self, embedder: Arc<dyn Embedder>, id: impl Into<String>) -> Self {
        self.embedder = Some((embedder, id.into()));
        self
    }

    /// Disable the Merkle-chain provenance recorder for this config.
    pub fn without_provenance(mut self) -> Self {
        self.provenance = false;
        self
    }

    /// Disable the episode-to-graph projector for this config.
    pub fn without_projector(mut self) -> Self {
        self.projector = false;
        self
    }

    /// Use the built-in extractor only; skip the LLM gap-filling pass.
    pub fn builtin_extractor_only(mut self) -> Self {
        self.extractor_mode = ExtractorMode::BuiltinOnly;
        self
    }

    /// Delegate entity extraction entirely to `extractor`.
    pub fn custom_extractor(mut self, extractor: Arc<dyn EntityExtractor>) -> Self {
        self.extractor_mode = ExtractorMode::Custom(extractor);
        self
    }

    /// Override the scope used to store graph projections.
    pub fn projection_scope(mut self, scope: Scope) -> Self {
        self.projection_scope = scope;
        self
    }

    /// Override the provenance actor label recorded in the audit chain.
    /// Defaults to a generic framework label; set this to distinguish
    /// App instances in a shared audit store.
    pub fn provenance_actor(mut self, actor: impl Into<String>) -> Self {
        self.provenance_actor = actor.into();
        self
    }

    /// Minimum graph hits needed to include graph context in a recall
    /// response. Defaults to `1`.
    pub fn min_graph_hits(mut self, n: usize) -> Self {
        self.min_graph_hits = n;
        self
    }
}

/// The assembled GraphRAG stack produced by [`GraphRagConfig::wire`].
///
/// `build()` installs [`long_term`][Self::long_term] as the App's
/// `LongTermMemory` port; [`graph_aware`][Self::graph_aware] is the same
/// port held concrete, for callers that need `GraphAwareLongTerm` directly.
/// The remaining fields are the shared handles the episode-to-graph
/// projector reuses: the same `graph`, `extractor`, `projection_scope`, and
/// `metrics` the long-term port was composed from.
pub(crate) struct WiredGraphRag {
    pub long_term: Arc<dyn LongTermMemory>,
    /// The same long-term port as `long_term`, held as its concrete type so
    /// callers that need `GraphAwareLongTerm`'s `recall_traced` capability
    /// (the ops-viz graph explorer) don't have to downcast a trait object.
    pub graph_aware: Arc<GraphAwareLongTerm>,
    pub metrics: Arc<RecallMetrics>,
    pub graph: Arc<dyn KnowledgeGraph>,
    pub extractor: Arc<dyn EntityExtractor>,
    pub projection_scope: Scope,
    pub projector_enabled: bool,
}

/// The vector store and raw knowledge graph for one backend.
struct VectorAndGraph {
    vector: Arc<dyn FilterableLongTermMemory>,
    raw_graph: Arc<dyn KnowledgeGraph>,
}

/// Inputs that decide whether and how the provenance recorder wraps the
/// raw knowledge graph.
struct ProvenanceWiring {
    enabled: bool,
    is_remote: bool,
    actor: String,
    sqlite_path: Option<std::path::PathBuf>,
}

impl GraphRagConfig {
    /// Assemble the whole GraphRAG stack from this config.
    ///
    /// `llm` is the App's mandatory LLM client (used by the LLM
    /// gap-filling extractor); `sqlite_path` is the App's persistent
    /// sqlite path, reused for the on-disk provenance database on the
    /// remote backend.
    pub(crate) async fn wire(
        self,
        llm: Arc<dyn LlmClient>,
        sqlite_path: Option<std::path::PathBuf>,
    ) -> Result<WiredGraphRag, Error> {
        let (embedder, embedder_id) = self.resolve_embedder()?;
        let provenance = ProvenanceWiring {
            enabled: self.provenance,
            is_remote: matches!(self.backend, Backend::Remote { .. }),
            actor: self.provenance_actor,
            sqlite_path,
        };
        let parts = build_vector_and_graph(self.backend, embedder, embedder_id).await?;
        let graph = wrap_provenance(parts.raw_graph, provenance)?;
        let extractor = build_extractor(self.extractor_mode, llm);

        let metrics = Arc::new(RecallMetrics::default());
        let graph_aware = Arc::new(
            GraphAwareLongTerm::builder()
                .extractor(extractor.clone())
                .metrics(metrics.clone())
                .min_graph_hits(self.min_graph_hits)
                .build(parts.vector, graph.clone()),
        );
        let long_term: Arc<dyn LongTermMemory> = graph_aware.clone();

        Ok(WiredGraphRag {
            long_term,
            graph_aware,
            metrics,
            graph,
            extractor,
            projection_scope: self.projection_scope,
            projector_enabled: self.projector,
        })
    }

    /// The configured embedder + its identifier, defaulting to fastembed.
    fn resolve_embedder(&self) -> Result<(Arc<dyn Embedder>, String), Error> {
        if let Some((embedder, id)) = &self.embedder {
            return Ok((embedder.clone(), id.clone()));
        }
        let fastembed = klieo_embed_common::FastEmbedEmbedder::new()?;
        Ok((Arc::new(fastembed), FASTEMBED_ID.to_string()))
    }
}

/// Construct the per-backend vector store and raw (unwrapped) graph.
async fn build_vector_and_graph(
    backend: Backend,
    embedder: Arc<dyn Embedder>,
    embedder_id: String,
) -> Result<VectorAndGraph, Error> {
    match backend {
        Backend::InMemory => Ok(VectorAndGraph {
            vector: Arc::new(InMemoryFilterableLongTerm::new(embedder, embedder_id)),
            raw_graph: Arc::new(klieo_memory_graph::InMemoryGraph::default()),
        }),
        Backend::Remote { qdrant_url, neo4j } => {
            let cfg =
                klieo_memory_qdrant::QdrantConfig::new(qdrant_url).with_embedder_id(embedder_id);
            let qdrant = klieo_memory_qdrant::MemoryQdrant::new(cfg, embedder).await?;
            let neo = klieo_memory_neo4j::MemoryNeo4j::new(neo4j).await?;
            let graph = klieo_memory_graph_neo4j::Neo4jKnowledgeGraph::new(neo.neo4j_handle());
            Ok(VectorAndGraph {
                vector: qdrant.qdrant_long_term,
                raw_graph: Arc::new(graph),
            })
        }
    }
}

/// Wrap `raw_graph` with the provenance recorder when enabled. The
/// remote backend persists provenance to the App's sqlite path (or a
/// default file); the in-memory backend uses an ephemeral `:memory:`
/// database. The configured actor is recorded on every emitted event.
fn wrap_provenance(
    raw_graph: Arc<dyn KnowledgeGraph>,
    wiring: ProvenanceWiring,
) -> Result<Arc<dyn KnowledgeGraph>, Error> {
    if !wiring.enabled {
        return Ok(raw_graph);
    }
    let repo = if wiring.is_remote {
        let path = wiring
            .sqlite_path
            .unwrap_or_else(|| std::path::PathBuf::from(DEFAULT_PROVENANCE_DB_PATH));
        klieo_provenance::SqliteProvenanceRepository::open(path)
    } else {
        klieo_provenance::SqliteProvenanceRepository::open_in_memory()
    }
    .map_err(|e| Error::Downstream(Box::new(e)))?;
    Ok(Arc::new(ProvenanceKnowledgeGraph::new(
        raw_graph,
        Arc::new(repo),
        wiring.actor,
    )))
}

/// Select the entity extractor for the configured mode. The LLM
/// gap-filling fallback uses the App's mandatory `llm` client.
fn build_extractor(mode: ExtractorMode, llm: Arc<dyn LlmClient>) -> Arc<dyn EntityExtractor> {
    match mode {
        ExtractorMode::Custom(extractor) => extractor,
        ExtractorMode::BuiltinOnly => Arc::new(BuiltinExtractor::default()),
        ExtractorMode::BuiltinThenLlm => Arc::new(FallbackExtractor::new(
            BuiltinExtractor::default(),
            LlmEntityExtractor::new(llm),
        )),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, Ordering};

    #[test]
    fn in_memory_defaults_all_on() {
        let c = GraphRagConfig::in_memory();
        assert!(c.provenance);
        assert!(c.projector);
        assert!(matches!(c.extractor_mode, ExtractorMode::BuiltinThenLlm));
        assert_eq!(c.min_graph_hits, 1);
    }

    #[test]
    fn opt_outs_flip_flags() {
        let c = GraphRagConfig::in_memory()
            .without_provenance()
            .without_projector()
            .builtin_extractor_only();
        assert!(!c.provenance);
        assert!(!c.projector);
        assert!(matches!(c.extractor_mode, ExtractorMode::BuiltinOnly));
    }

    // `remote` is constructed with the re-exported `Neo4jConfig` (proving the
    // facade re-export resolves) and must store the qdrant URL in the
    // Remote backend variant.
    #[test]
    fn remote_stores_qdrant_url() {
        let neo4j = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "password".to_string());
        let c = GraphRagConfig::remote("http://localhost:6333", neo4j);
        match c.backend {
            Backend::Remote { qdrant_url, .. } => {
                assert_eq!(qdrant_url, "http://localhost:6333");
            }
            Backend::InMemory => panic!("remote() must produce Backend::Remote"),
        }
    }

    const TEST_EMBEDDER_DIM: usize = 384;

    fn fake_llm() -> Arc<dyn LlmClient> {
        Arc::new(klieo_core::test_utils::FakeLlmClient::new("fake"))
    }

    /// Deterministic, non-zero-vector embedder so cosine ranking is
    /// actually exercised (all-zero vectors would tie every fact).
    fn test_embedder() -> Arc<klieo_embed_common::FakeEmbedder> {
        Arc::new(klieo_embed_common::FakeEmbedder::new(TEST_EMBEDDER_DIM))
    }

    #[test]
    fn provenance_actor_defaults_and_overrides() {
        let default = GraphRagConfig::in_memory();
        assert_eq!(default.provenance_actor, PROVENANCE_ACTOR);
        let custom = GraphRagConfig::in_memory().provenance_actor("tenant-a");
        assert_eq!(custom.provenance_actor, "tenant-a");
    }

    #[tokio::test]
    async fn in_memory_wire_produces_working_long_term() {
        let cfg = GraphRagConfig::in_memory().embedder(test_embedder(), "test-fake");
        let wired = cfg
            .wire(fake_llm(), None)
            .await
            .expect("in-memory wire should not touch any external service");
        assert!(wired.projector_enabled, "projector defaults on");

        let scope = Scope::Workspace("wire-test".to_string());
        let target = "the wired port round-trips facts";
        wired
            .long_term
            .remember(scope.clone(), klieo_core::memory::Fact::new(target))
            .await
            .expect("remember through the wired long-term port");
        wired
            .long_term
            .remember(
                scope.clone(),
                klieo_core::memory::Fact::new("unrelated weather report"),
            )
            .await
            .expect("remember of a distractor fact");
        let hits = wired
            .long_term
            .recall(scope, "the wired port round-trips facts", 5)
            .await
            .expect("recall through the wired long-term port");
        assert_eq!(
            hits.first().map(|f| f.text.as_str()),
            Some(target),
            "the query-matching fact must rank first under real cosine ranking; got {hits:?}"
        );
    }

    #[tokio::test]
    async fn wire_without_provenance_skips_recorder() {
        let cfg = GraphRagConfig::in_memory()
            .embedder(test_embedder(), "test-fake")
            .without_provenance();
        let wired = cfg
            .wire(fake_llm(), None)
            .await
            .expect("wire should succeed with provenance disabled");
        let scope = Scope::Workspace("wire-no-prov".to_string());
        wired
            .long_term
            .remember(
                scope.clone(),
                klieo_core::memory::Fact::new("no provenance here"),
            )
            .await
            .expect("remember should still work without provenance");
        let hits = wired
            .long_term
            .recall(scope, "no provenance here", 5)
            .await
            .expect("recall should still work without provenance");
        assert!(hits.iter().any(|f| f.text.contains("no provenance")));
    }

    /// Sentinel extractor that records whether `extract` was called.
    struct SentinelExtractor(Arc<AtomicBool>);

    #[async_trait::async_trait]
    impl klieo_memory_graph::EntityExtractor for SentinelExtractor {
        async fn extract(
            &self,
            _text: &str,
            _hints: &[klieo_memory_graph::EntityRef],
        ) -> Result<Vec<klieo_memory_graph::EntityRef>, klieo_core::error::MemoryError> {
            self.0.store(true, Ordering::Release);
            Ok(vec![])
        }
    }

    #[tokio::test]
    async fn custom_extractor_is_called_on_remember() {
        let called = Arc::new(AtomicBool::new(false));
        let sentinel = Arc::new(SentinelExtractor(called.clone()));
        let cfg = GraphRagConfig::in_memory()
            .embedder(test_embedder(), "test-fake")
            .custom_extractor(sentinel);
        let wired = cfg
            .wire(fake_llm(), None)
            .await
            .expect("in-memory wire with custom extractor should succeed");
        let scope = Scope::Workspace("custom-extractor-test".to_string());
        wired
            .long_term
            .remember(scope, klieo_core::memory::Fact::new("sentinel text"))
            .await
            .expect("remember should invoke the extractor");
        assert!(
            called.load(Ordering::Acquire),
            "custom_extractor must be called during remember"
        );
    }
}