Skip to main content

klieo_memory_graph_rag/
import.rs

1//! `KnowledgeIngestionPipeline` + `ImportSource` — bulk-import external
2//! documents into a `LongTermMemory`.
3//!
4//! Extension point for future connectors (Confluence, Notion, Jira, DORA
5//! regulatory texts). Connectors implement `ImportSource` in their own
6//! crates; `KnowledgeIngestionPipeline::run` wraps each fetched document
7//! as a `Fact`, threading `entity_hints` and `valid_from` through
8//! `Fact::metadata` so downstream `GraphAwareLongTerm::remember` (M2)
9//! can recover them on the index path.
10
11use crate::metadata::{
12    KEY_ENTITIES, KEY_IMPORT_SOURCE_ID, KEY_TITLE, KEY_VALID_FROM, RESERVED_METADATA_KEYS,
13};
14use async_trait::async_trait;
15use chrono::{DateTime, Utc};
16use klieo_core::error::MemoryError;
17use klieo_core::memory::{Fact, LongTermMemory, Scope};
18use klieo_memory_graph::EntityRef;
19use std::sync::Arc;
20
21/// One document fetched from an `ImportSource`.
22///
23/// Wrapped as a `Fact` by `KnowledgeIngestionPipeline::run`: `body`
24/// becomes `Fact.text`; `entity_hints`, `valid_from`, `id`, and `title`
25/// land under reserved keys in `Fact.metadata`; caller-supplied
26/// `metadata` is merged (pipeline writes reserved keys LAST, so they
27/// overwrite any collision — callers MUST NOT use reserved key names).
28#[derive(Debug, Clone)]
29#[non_exhaustive]
30pub struct ImportDocument {
31    /// Stable external identifier (Confluence page id, Jira key, file
32    /// hash, …). Persisted as `Fact.metadata.import_source_id`.
33    pub id: String,
34    /// Display name. Persisted as `Fact.metadata.title`.
35    pub title: String,
36    /// Document body — becomes `Fact.text`.
37    pub body: String,
38    /// Caller-supplied opaque metadata. Merged with pipeline-reserved
39    /// keys (`valid_from`, `entities`, `import_source_id`, `title`).
40    pub metadata: serde_json::Value,
41    /// Pre-known entities (`Article`, `Obligation`, etc. for DORA;
42    /// owners for Confluence pages, …). Persisted as
43    /// `Fact.metadata.entities` so `GraphAwareLongTerm::remember`
44    /// surfaces them via the extractor's hint path.
45    pub entity_hints: Vec<EntityRef>,
46    /// When the document is valid in the world (regulation effective
47    /// date, contract start, …). Persisted as `Fact.metadata.valid_from`
48    /// in RFC-3339; `GraphAwareLongTerm` parses it back.
49    pub valid_from: Option<DateTime<Utc>>,
50}
51
52impl ImportDocument {
53    /// Build a document from required fields. Optional fields use their
54    /// `Default` values; set them via the builder-style `with_*` methods.
55    pub fn new(id: impl Into<String>, title: impl Into<String>, body: impl Into<String>) -> Self {
56        Self {
57            id: id.into(),
58            title: title.into(),
59            body: body.into(),
60            metadata: serde_json::Value::Null,
61            entity_hints: Vec::new(),
62            valid_from: None,
63        }
64    }
65
66    /// Attach caller-supplied metadata (merged into `Fact.metadata` by
67    /// `KnowledgeIngestionPipeline::run`).
68    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
69        self.metadata = metadata;
70        self
71    }
72
73    /// Pre-seed entities the extractor should treat as authoritative.
74    pub fn with_entity_hints(mut self, hints: Vec<EntityRef>) -> Self {
75        self.entity_hints = hints;
76        self
77    }
78
79    /// Bi-temporal `valid_from` (e.g. regulation effective date).
80    pub fn with_valid_from(mut self, valid_from: DateTime<Utc>) -> Self {
81        self.valid_from = Some(valid_from);
82        self
83    }
84}
85
86/// Errors raised by `KnowledgeIngestionPipeline::run` and `ImportSource::fetch`.
87///
88/// `Fetch` carries the upstream's failure shape verbatim (already
89/// redacted by the connector); `Memory` lifts a `MemoryError` from the
90/// underlying `LongTermMemory` so the pipeline error type unifies both
91/// halves of the ingest path.
92#[derive(Debug, thiserror::Error)]
93#[non_exhaustive]
94pub enum ImportError {
95    /// Source-side fetch failed (HTTP error, auth failure, file not
96    /// found, …). Connectors are responsible for redacting secrets
97    /// from the message before propagating.
98    #[error("import source fetch failed: {0}")]
99    Fetch(String),
100
101    /// Downstream memory write failed.
102    #[error("import memory write failed: {0}")]
103    Memory(#[from] MemoryError),
104}
105
106/// External knowledge source — implementors live in their own crates
107/// (e.g. `klieo-import-confluence`, `klieo-tools-regulation` for M6).
108#[async_trait]
109pub trait ImportSource: Send + Sync {
110    /// Fetch all documents this source exposes. Connectors decide
111    /// pagination, batching, and error-vs-empty semantics; the pipeline
112    /// consumes the returned `Vec` once per `run` invocation.
113    async fn fetch(&self) -> Result<Vec<ImportDocument>, ImportError>;
114}
115
116/// Bulk-import pipeline that pulls `ImportDocument`s from an
117/// `ImportSource` and writes them as `Fact`s into a `LongTermMemory`.
118pub struct KnowledgeIngestionPipeline {
119    source: Arc<dyn ImportSource>,
120    memory: Arc<dyn LongTermMemory>,
121    scope: Scope,
122}
123
124impl KnowledgeIngestionPipeline {
125    /// Build with the source, destination memory, and target scope.
126    pub fn new(
127        source: Arc<dyn ImportSource>,
128        memory: Arc<dyn LongTermMemory>,
129        scope: Scope,
130    ) -> Self {
131        Self {
132            source,
133            memory,
134            scope,
135        }
136    }
137
138    /// Fetch all documents from `source` and persist them under `scope`.
139    ///
140    /// Fail-fast: the first per-document `memory.remember` failure aborts
141    /// the batch with `ImportError::Memory`. Documents persisted before the
142    /// failure remain in the memory backend — there is no rollback.
143    /// On full success returns the count of persisted documents.
144    pub async fn run(&self) -> Result<usize, ImportError> {
145        let docs = self.source.fetch().await?;
146        let mut persisted = 0;
147        for doc in docs {
148            let fact = self.build_fact(doc)?;
149            self.memory.remember(self.scope.clone(), fact).await?;
150            persisted += 1;
151        }
152        Ok(persisted)
153    }
154
155    fn build_fact(&self, doc: ImportDocument) -> Result<Fact, ImportError> {
156        let mut map = match doc.metadata {
157            serde_json::Value::Object(m) => m,
158            _ => serde_json::Map::new(),
159        };
160
161        if let Some(vf) = doc.valid_from {
162            map.insert(
163                KEY_VALID_FROM.into(),
164                serde_json::Value::String(vf.to_rfc3339()),
165            );
166        }
167        if !doc.entity_hints.is_empty() {
168            let value = serde_json::to_value(&doc.entity_hints)
169                .map_err(|e| ImportError::Memory(MemoryError::Serialization(e.to_string())))?;
170            map.insert(KEY_ENTITIES.into(), value);
171        }
172        map.insert(
173            KEY_IMPORT_SOURCE_ID.into(),
174            serde_json::Value::String(doc.id),
175        );
176        map.insert(KEY_TITLE.into(), serde_json::Value::String(doc.title));
177
178        Ok(Fact::new(doc.body).with_metadata(serde_json::Value::Object(map)))
179    }
180}
181
182/// Reserved-key audit — kept as a `pub fn` so M6's connector test
183/// harness can guard against callers stomping reserved keys.
184#[doc(hidden)]
185pub fn reserved_metadata_keys() -> &'static [&'static str] {
186    RESERVED_METADATA_KEYS
187}