klieo-memory-graph-rag 3.4.0

//! [`RegulationIngestTool`] — [`ImportSource`] for structured
//! regulatory texts (DORA, AI-Act, …).
//!
//! Each [`RegulationDocument`] becomes one [`ImportDocument`] whose
//! `entity_hints` are pre-seeded with the regulation's `Article`,
//! `Obligation`, and `Deadline` entities so
//! `GraphAwareLongTerm::remember` surfaces them on the graph-index
//! path. `valid_from` defaults to
//! [`klieo_memory_graph::known_epoch`] so
//! the recorded knowledge is treated as all-time valid; callers
//! that need stricter bi-temporal semantics (e.g. an article that
//! enters force in the future) override via
//! [`RegulationDocument::with_valid_from`].
//!
//! The tool itself holds an in-memory `Vec<RegulationDocument>`
//! supplied at construction. Production deployments will replace
//! the in-memory store with a real parser (XML / Eur-Lex / …) but
//! the ImportSource contract stays identical — only `fetch()`'s
//! source changes.

use crate::import::{ImportDocument, ImportError, ImportSource};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_memory_graph::{known_epoch, EntityRef, EntityType};

/// One regulatory text — the unit `RegulationIngestTool` ingests.
///
/// `entity_hints` is derived from `article_ids`, `obligations`,
/// and `deadlines` at construction time so callers can't forget
/// to wire one of them. Build via [`Self::new`] +
/// the `with_*` builders.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct RegulationDocument {
    /// Stable external identifier (e.g. `"dora-art-5"`,
    /// `"eu-ai-act-art-13"`).
    pub id: String,
    /// Display title (e.g. `"DORA Article 5 — ICT risk
    /// management"`).
    pub title: String,
    /// Full regulatory text — becomes `Fact.text` after ingestion.
    pub body: String,
    /// Article identifiers this document encodes
    /// (e.g. `["DORA-Art-5", "DORA-Art-6"]`). Each becomes an
    /// `Article` entity hint.
    pub article_ids: Vec<String>,
    /// Obligation identifiers this document creates
    /// (e.g. `["ICT-Risk-Management"]`). Each becomes an
    /// `Obligation` entity hint.
    pub obligations: Vec<String>,
    /// Deadline identifiers this document imposes
    /// (e.g. `["DORA-Compliance-2026-01-17"]`). Each becomes a
    /// `Deadline` entity hint.
    pub deadlines: Vec<String>,
    /// When this regulation becomes valid in the world. Defaults
    /// to [`klieo_memory_graph::known_epoch`] (all-time) when
    /// unset; override for future-dated articles via
    /// [`Self::with_valid_from`].
    pub valid_from: Option<DateTime<Utc>>,
}

impl RegulationDocument {
    /// Article / obligation / deadline lists default to empty;
    /// add them via the `with_*` builders.
    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(),
            article_ids: Vec::new(),
            obligations: Vec::new(),
            deadlines: Vec::new(),
            valid_from: None,
        }
    }

    /// Append article identifiers. Chains additively — calling
    /// `with_articles` twice keeps both batches.
    pub fn with_articles<I, S>(mut self, ids: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.article_ids.extend(ids.into_iter().map(Into::into));
        self
    }

    /// Append obligation identifiers. Chains additively.
    pub fn with_obligations<I, S>(mut self, ids: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.obligations.extend(ids.into_iter().map(Into::into));
        self
    }

    /// Append deadline identifiers. Chains additively.
    pub fn with_deadlines<I, S>(mut self, ids: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.deadlines.extend(ids.into_iter().map(Into::into));
        self
    }

    /// Override the default `valid_from`
    /// ([`klieo_memory_graph::known_epoch`]) for future-dated
    /// articles.
    pub fn with_valid_from(mut self, valid_from: DateTime<Utc>) -> Self {
        self.valid_from = Some(valid_from);
        self
    }

    fn entity_hints(&self) -> Vec<EntityRef> {
        let mut hints = Vec::with_capacity(
            self.article_ids.len() + self.obligations.len() + self.deadlines.len(),
        );
        for id in &self.article_ids {
            hints.push(EntityRef::new(EntityType::Article, id.as_str()));
        }
        for id in &self.obligations {
            hints.push(EntityRef::new(EntityType::Obligation, id.as_str()));
        }
        for id in &self.deadlines {
            hints.push(EntityRef::new(EntityType::Deadline, id.as_str()));
        }
        hints
    }
}

/// [`ImportSource`] that ingests structured regulatory texts.
///
/// Construction takes the `Vec<RegulationDocument>` directly so
/// tests + demos can stage fixtures inline. A future production
/// connector will read DORA / AI-Act / GDPR XML feeds and produce
/// the same `Vec` shape — the `ImportSource` contract is
/// unchanged.
#[non_exhaustive]
pub struct RegulationIngestTool {
    documents: Vec<RegulationDocument>,
}

impl RegulationIngestTool {
    /// Build with the supplied corpus. The `Vec` is consumed
    /// once per [`Self::fetch`] call — re-running the pipeline
    /// re-emits the same documents (idempotent on the graph-index
    /// path because `GraphAwareLongTerm::remember` deduplicates
    /// by content).
    pub fn new(documents: Vec<RegulationDocument>) -> Self {
        Self { documents }
    }
}

#[async_trait]
impl ImportSource for RegulationIngestTool {
    async fn fetch(&self) -> Result<Vec<ImportDocument>, ImportError> {
        let docs = self
            .documents
            .iter()
            .map(|reg| {
                let valid_from = reg.valid_from.unwrap_or_else(known_epoch);
                ImportDocument::new(reg.id.clone(), reg.title.clone(), reg.body.clone())
                    .with_entity_hints(reg.entity_hints())
                    .with_valid_from(valid_from)
            })
            .collect();
        Ok(docs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    fn dora_article_5() -> RegulationDocument {
        RegulationDocument::new(
            "dora-art-5",
            "DORA Article 5 — ICT risk management",
            "Financial entities shall have in place a sound, comprehensive and \
             well-documented ICT risk management framework.",
        )
        .with_articles(["DORA-Art-5"])
        .with_obligations(["ICT-Risk-Management"])
        .with_deadlines(["DORA-Compliance-2026-01-17"])
    }

    #[tokio::test]
    async fn fetch_emits_one_import_document_per_regulation() {
        let tool = RegulationIngestTool::new(vec![dora_article_5()]);
        let docs = tool.fetch().await.unwrap();
        assert_eq!(docs.len(), 1);
        assert_eq!(docs[0].id, "dora-art-5");
        assert_eq!(docs[0].title, "DORA Article 5 — ICT risk management");
    }

    #[tokio::test]
    async fn fetch_seeds_entity_hints_with_typed_regulatory_variants() {
        let tool = RegulationIngestTool::new(vec![dora_article_5()]);
        let docs = tool.fetch().await.unwrap();
        let hints = &docs[0].entity_hints;
        let kinds: Vec<&EntityType> = hints.iter().map(|h| &h.entity_type).collect();
        assert!(kinds.contains(&&EntityType::Article));
        assert!(kinds.contains(&&EntityType::Obligation));
        assert!(kinds.contains(&&EntityType::Deadline));
        // Names lowercased by EntityRef::new.
        let names: Vec<&str> = hints.iter().map(|h| h.name.as_str()).collect();
        assert!(names.contains(&"dora-art-5"));
        assert!(names.contains(&"ict-risk-management"));
        assert!(names.contains(&"dora-compliance-2026-01-17"));
    }

    #[tokio::test]
    async fn fetch_defaults_valid_from_to_known_epoch() {
        let tool = RegulationIngestTool::new(vec![dora_article_5()]);
        let docs = tool.fetch().await.unwrap();
        assert_eq!(docs[0].valid_from, Some(known_epoch()));
    }

    #[tokio::test]
    async fn fetch_honours_explicit_valid_from_override() {
        let future = Utc.with_ymd_and_hms(2027, 1, 17, 0, 0, 0).single().unwrap();
        let tool = RegulationIngestTool::new(vec![dora_article_5().with_valid_from(future)]);
        let docs = tool.fetch().await.unwrap();
        assert_eq!(docs[0].valid_from, Some(future));
    }

    #[tokio::test]
    async fn empty_regulation_list_yields_empty_fetch() {
        let tool = RegulationIngestTool::new(Vec::new());
        let docs = tool.fetch().await.unwrap();
        assert!(docs.is_empty());
    }

    #[test]
    fn entity_hints_emit_in_article_obligation_deadline_order() {
        let doc = dora_article_5();
        let hints = doc.entity_hints();
        // Articles first, then obligations, then deadlines —
        // ordering is part of the contract so downstream graph
        // CO_OCCURS edges are deterministic across runs.
        assert_eq!(hints[0].entity_type, EntityType::Article);
        assert_eq!(hints[1].entity_type, EntityType::Obligation);
        assert_eq!(hints[2].entity_type, EntityType::Deadline);
    }

    #[test]
    fn entity_hints_handles_multi_value_lists() {
        let doc = RegulationDocument::new("multi", "Multi", "body")
            .with_articles(["A1", "A2"])
            .with_obligations(["O1"])
            .with_deadlines(["D1", "D2", "D3"]);
        let hints = doc.entity_hints();
        assert_eq!(hints.len(), 6);
        // Articles first (×2), then obligations (×1), then deadlines (×3).
        assert_eq!(hints[0].entity_type, EntityType::Article);
        assert_eq!(hints[0].name, "a1");
        assert_eq!(hints[1].entity_type, EntityType::Article);
        assert_eq!(hints[1].name, "a2");
        assert_eq!(hints[2].entity_type, EntityType::Obligation);
        assert_eq!(hints[2].name, "o1");
        assert_eq!(hints[3].entity_type, EntityType::Deadline);
        assert_eq!(hints[3].name, "d1");
        assert_eq!(hints[4].entity_type, EntityType::Deadline);
        assert_eq!(hints[5].entity_type, EntityType::Deadline);
    }

    #[test]
    fn with_articles_chains_additively_across_calls() {
        // Pre-0.8 builders REPLACED on each call — chaining
        // `.with_articles(["A1"]).with_articles(["A2"])` silently
        // dropped A1. Post-0.8 they extend: both survive.
        let doc = RegulationDocument::new("multi-call", "Multi", "body")
            .with_articles(["A1"])
            .with_articles(["A2"])
            .with_obligations(["O1"])
            .with_obligations(["O2"])
            .with_deadlines(["D1"])
            .with_deadlines(["D2", "D3"]);
        assert_eq!(
            doc.article_ids,
            vec!["A1".to_string(), "A2".to_string()],
            "chained with_articles must append, not replace"
        );
        assert_eq!(doc.obligations, vec!["O1".to_string(), "O2".to_string()]);
        assert_eq!(
            doc.deadlines,
            vec!["D1".to_string(), "D2".to_string(), "D3".to_string()]
        );
    }
}