klieo-memory-graph-rag 2.1.0

//! `KnowledgeIngestionPipeline` + `ImportSource` — bulk-import external
//! documents into a `LongTermMemory`.
//!
//! Extension point for future connectors (Confluence, Notion, Jira, DORA
//! regulatory texts). Connectors implement `ImportSource` in their own
//! crates; `KnowledgeIngestionPipeline::run` wraps each fetched document
//! as a `Fact`, threading `entity_hints` and `valid_from` through
//! `Fact::metadata` so downstream `GraphAwareLongTerm::remember` (M2)
//! can recover them on the index path.

use crate::metadata::{
    KEY_ENTITIES, KEY_IMPORT_SOURCE_ID, KEY_TITLE, KEY_VALID_FROM, RESERVED_METADATA_KEYS,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_core::error::MemoryError;
use klieo_core::memory::{Fact, LongTermMemory, Scope};
use klieo_memory_graph::EntityRef;
use std::sync::Arc;

/// One document fetched from an `ImportSource`.
///
/// Wrapped as a `Fact` by `KnowledgeIngestionPipeline::run`: `body`
/// becomes `Fact.text`; `entity_hints`, `valid_from`, `id`, and `title`
/// land under reserved keys in `Fact.metadata`; caller-supplied
/// `metadata` is merged (pipeline writes reserved keys LAST, so they
/// overwrite any collision — callers MUST NOT use reserved key names).
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ImportDocument {
    /// Stable external identifier (Confluence page id, Jira key, file
    /// hash, …). Persisted as `Fact.metadata.import_source_id`.
    pub id: String,
    /// Display name. Persisted as `Fact.metadata.title`.
    pub title: String,
    /// Document body — becomes `Fact.text`.
    pub body: String,
    /// Caller-supplied opaque metadata. Merged with pipeline-reserved
    /// keys (`valid_from`, `entities`, `import_source_id`, `title`).
    pub metadata: serde_json::Value,
    /// Pre-known entities (`Article`, `Obligation`, etc. for DORA;
    /// owners for Confluence pages, …). Persisted as
    /// `Fact.metadata.entities` so `GraphAwareLongTerm::remember`
    /// surfaces them via the extractor's hint path.
    pub entity_hints: Vec<EntityRef>,
    /// When the document is valid in the world (regulation effective
    /// date, contract start, …). Persisted as `Fact.metadata.valid_from`
    /// in RFC-3339; `GraphAwareLongTerm` parses it back.
    pub valid_from: Option<DateTime<Utc>>,
}

impl ImportDocument {
    /// Build a document from required fields. Optional fields use their
    /// `Default` values; set them via the builder-style `with_*` methods.
    pub fn new(id: impl Into<String>, title: impl Into<String>, body: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            title: title.into(),
            body: body.into(),
            metadata: serde_json::Value::Null,
            entity_hints: Vec::new(),
            valid_from: None,
        }
    }

    /// Attach caller-supplied metadata (merged into `Fact.metadata` by
    /// `KnowledgeIngestionPipeline::run`).
    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
        self.metadata = metadata;
        self
    }

    /// Pre-seed entities the extractor should treat as authoritative.
    pub fn with_entity_hints(mut self, hints: Vec<EntityRef>) -> Self {
        self.entity_hints = hints;
        self
    }

    /// Bi-temporal `valid_from` (e.g. regulation effective date).
    pub fn with_valid_from(mut self, valid_from: DateTime<Utc>) -> Self {
        self.valid_from = Some(valid_from);
        self
    }
}

/// Errors raised by `KnowledgeIngestionPipeline::run` and `ImportSource::fetch`.
///
/// `Fetch` carries the upstream's failure shape verbatim (already
/// redacted by the connector); `Memory` lifts a `MemoryError` from the
/// underlying `LongTermMemory` so the pipeline error type unifies both
/// halves of the ingest path.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ImportError {
    /// Source-side fetch failed (HTTP error, auth failure, file not
    /// found, …). Connectors are responsible for redacting secrets
    /// from the message before propagating.
    #[error("import source fetch failed: {0}")]
    Fetch(String),

    /// Downstream memory write failed.
    #[error("import memory write failed: {0}")]
    Memory(#[from] MemoryError),
}

/// External knowledge source — implementors live in their own crates
/// (e.g. `klieo-import-confluence`, `klieo-tools-regulation` for M6).
#[async_trait]
pub trait ImportSource: Send + Sync {
    /// Fetch all documents this source exposes. Connectors decide
    /// pagination, batching, and error-vs-empty semantics; the pipeline
    /// consumes the returned `Vec` once per `run` invocation.
    async fn fetch(&self) -> Result<Vec<ImportDocument>, ImportError>;
}

/// Bulk-import pipeline that pulls `ImportDocument`s from an
/// `ImportSource` and writes them as `Fact`s into a `LongTermMemory`.
pub struct KnowledgeIngestionPipeline {
    source: Arc<dyn ImportSource>,
    memory: Arc<dyn LongTermMemory>,
    scope: Scope,
}

impl KnowledgeIngestionPipeline {
    /// Build with the source, destination memory, and target scope.
    pub fn new(
        source: Arc<dyn ImportSource>,
        memory: Arc<dyn LongTermMemory>,
        scope: Scope,
    ) -> Self {
        Self {
            source,
            memory,
            scope,
        }
    }

    /// Fetch all documents from `source` and persist them under `scope`.
    ///
    /// Fail-fast: the first per-document `memory.remember` failure aborts
    /// the batch with `ImportError::Memory`. Documents persisted before the
    /// failure remain in the memory backend — there is no rollback.
    /// On full success returns the count of persisted documents.
    pub async fn run(&self) -> Result<usize, ImportError> {
        let docs = self.source.fetch().await?;
        let mut persisted = 0;
        for doc in docs {
            let fact = self.build_fact(doc)?;
            self.memory.remember(self.scope.clone(), fact).await?;
            persisted += 1;
        }
        Ok(persisted)
    }

    fn build_fact(&self, doc: ImportDocument) -> Result<Fact, ImportError> {
        let mut map = match doc.metadata {
            serde_json::Value::Object(m) => m,
            _ => serde_json::Map::new(),
        };

        if let Some(vf) = doc.valid_from {
            map.insert(
                KEY_VALID_FROM.into(),
                serde_json::Value::String(vf.to_rfc3339()),
            );
        }
        if !doc.entity_hints.is_empty() {
            let value = serde_json::to_value(&doc.entity_hints)
                .map_err(|e| ImportError::Memory(MemoryError::Serialization(e.to_string())))?;
            map.insert(KEY_ENTITIES.into(), value);
        }
        map.insert(
            KEY_IMPORT_SOURCE_ID.into(),
            serde_json::Value::String(doc.id),
        );
        map.insert(KEY_TITLE.into(), serde_json::Value::String(doc.title));

        Ok(Fact::new(doc.body).with_metadata(serde_json::Value::Object(map)))
    }
}

/// Reserved-key audit — kept as a `pub fn` so M6's connector test
/// harness can guard against callers stomping reserved keys.
#[doc(hidden)]
pub fn reserved_metadata_keys() -> &'static [&'static str] {
    RESERVED_METADATA_KEYS
}