klieo_memory_graph_rag/

path_rag.rs

//! [`VerbalizationPipeline`] — LLM-driven prose summary of a
//! [`klieo_memory_graph::RetrievalPath`] under a
//! [`klieo_spec::QualityLoop`] refinement budget.
//!
//! Given a recall path produced by `ProvenanceKnowledgeGraph::recall_paths`,
//! the pipeline asks the LLM to compose a causal narrative — "X
//! co-occurs with Y, which is mentioned in fact F" — then iterates
//! through the [`VerbalizationCritic`] + [`LlmRefiner`] pair until
//! the critic accepts or the iteration budget is exhausted.
//!
//! The default critic passes when every `PathHop.fact_id` is
//! mentioned verbatim in the candidate prose; downstream extensions
//! can layer additional checks (e.g. regulatory-citation matching)
//! by composing their own [`klieo_spec::Critic`].

use async_trait::async_trait;
use klieo_core::llm::{ChatRequest, LlmClient, Message, Role};
use klieo_memory_graph::RetrievalPath;
use klieo_spec::{Critic, Critique, QualityLoop, QualityMetadata, Refiner, SpecError};
use std::sync::Arc;

/// Default iteration budget for the verbalisation quality loop.
/// Two passes — initial draft + one refinement — caps LLM spend per
/// verbalisation while leaving room for the critic to reject the
/// first attempt.
pub const DEFAULT_MAX_ITERATIONS: u8 = 2;

const SYSTEM_PROMPT: &str =
    "You are a retrieval-path verbaliser. Compose a concise causal narrative \
     that mentions every fact id supplied. Keep the prose plain and accurate; \
     never invent facts beyond the supplied list.";

/// LLM-driven verbaliser over a [`RetrievalPath`].
#[non_exhaustive]
pub struct VerbalizationPipeline {
    llm: Arc<dyn LlmClient>,
    max_iterations: u8,
}

impl VerbalizationPipeline {
    /// Construct the pipeline with the default
    /// [`DEFAULT_MAX_ITERATIONS`] budget. `llm` is shared across the
    /// initial draft and every refinement; the pipeline does not
    /// configure model selection — the supplied client handles
    /// model+temperature internally.
    pub fn new(llm: Arc<dyn LlmClient>) -> Self {
        Self {
            llm,
            max_iterations: DEFAULT_MAX_ITERATIONS,
        }
    }

    /// Override the per-verbalisation iteration cap. `0` is rejected
    /// at `verbalize` time with [`SpecError::Config`] (via the
    /// underlying [`QualityLoop`]).
    pub fn with_max_iterations(mut self, n: u8) -> Self {
        self.max_iterations = n;
        self
    }

    /// Verbalise `path` into a causal narrative. Returns the final
    /// text and the loop's run metadata (iteration count + whether
    /// the critic ultimately passed). Inspect
    /// [`QualityMetadata::passed`] before relying on the prose for
    /// downstream display.
    pub async fn verbalize(
        &self,
        path: &RetrievalPath,
    ) -> Result<(String, QualityMetadata), SpecError> {
        let hops_summary = render_hops(path);
        let initial = self.draft_initial(&hops_summary).await?;
        let critic = VerbalizationCritic::new(path);
        let refiner = LlmRefiner {
            llm: Arc::clone(&self.llm),
            hops_summary,
        };
        QualityLoop::new(critic, refiner)
            .with_max_iterations(self.max_iterations)
            .run(initial)
            .await
    }

    async fn draft_initial(&self, hops_summary: &str) -> Result<String, SpecError> {
        let req = ChatRequest::new(vec![
            Message {
                role: Role::System,
                content: SYSTEM_PROMPT.to_string(),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
            Message {
                role: Role::User,
                content: format!("Hops:\n{hops_summary}\n\nWrite the narrative."),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
        ]);
        complete(self.llm.as_ref(), req).await
    }
}

/// Critic that passes when the candidate verbalisation mentions
/// every `fact_id` from the supplied path. The check is verbatim
/// substring matching: the LLM must echo each `fact_id` once so the
/// downstream evidence-bundle linker can splice the prose back to
/// the chain entries it cites.
pub struct VerbalizationCritic {
    required_fact_ids: Vec<String>,
}

impl VerbalizationCritic {
    /// Snapshot the required fact ids from the supplied path so the
    /// critic owns its evaluation contract independent of the path's
    /// lifetime.
    pub fn new(path: &RetrievalPath) -> Self {
        let required_fact_ids = path.hops.iter().map(|h| h.fact_id.to_string()).collect();
        Self { required_fact_ids }
    }
}

#[async_trait]
impl Critic<String> for VerbalizationCritic {
    async fn evaluate(&self, candidate: &String, _iter: u8) -> Result<Critique, SpecError> {
        if self.required_fact_ids.is_empty() {
            return Ok(Critique::pass("no facts to verbalise"));
        }
        let missing: Vec<&str> = self
            .required_fact_ids
            .iter()
            .filter(|f| !candidate.contains(f.as_str()))
            .map(String::as_str)
            .collect();
        if missing.is_empty() {
            Ok(Critique::pass("every fact id is mentioned"))
        } else {
            Ok(Critique::fail(format!(
                "missing fact ids: {}",
                missing.join(", ")
            )))
        }
    }
}

/// Refiner that re-prompts the LLM with the critique feedback so
/// the next draft addresses the missing facts.
pub struct LlmRefiner {
    llm: Arc<dyn LlmClient>,
    hops_summary: String,
}

#[async_trait]
impl Refiner<String> for LlmRefiner {
    async fn refine(
        &self,
        current: String,
        critique: &Critique,
        _iter: u8,
    ) -> Result<String, SpecError> {
        let req = ChatRequest::new(vec![
            Message {
                role: Role::System,
                content: SYSTEM_PROMPT.to_string(),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
            Message {
                role: Role::User,
                content: format!(
                    "Hops:\n{}\n\nPrevious draft:\n{current}\n\nCritique:\n{}\n\nRevise the draft so the critique passes.",
                    self.hops_summary, critique.feedback,
                ),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
        ]);
        complete(self.llm.as_ref(), req).await
    }
}

fn render_hops(path: &RetrievalPath) -> String {
    let mut out = String::new();
    for (idx, hop) in path.hops.iter().enumerate() {
        out.push_str(&format!(
            "{}. entity={} fact_id={}",
            idx + 1,
            hop.entity.name,
            hop.fact_id
        ));
        if hop.chain_entry.is_some() {
            out.push_str(" provenance=signed");
        }
        out.push('\n');
    }
    out
}

async fn complete(llm: &dyn LlmClient, req: ChatRequest) -> Result<String, SpecError> {
    let resp = llm.complete(req).await.map_err(|e| SpecError::Downstream {
        message: "verbalisation LLM call failed".into(),
        source: Some(Box::new(e)),
    })?;
    Ok(resp.message.content)
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::ids::FactId;
    use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep};
    use klieo_memory_graph::{EntityRef, EntityType, PathHop};

    fn path_with_two_hops() -> RetrievalPath {
        RetrievalPath::with_hops(vec![
            PathHop::new(
                EntityRef::new(EntityType::Ticket, "T-1"),
                FactId::new("fact-alpha"),
            ),
            PathHop::new(
                EntityRef::new(EntityType::Policy, "P-1"),
                FactId::new("fact-beta"),
            ),
        ])
    }

    fn scripted(steps: Vec<&str>) -> Arc<FakeLlmClient> {
        Arc::new(
            FakeLlmClient::new("verbalisation-test").with_steps(
                steps
                    .into_iter()
                    .map(|s| FakeLlmStep::Text(s.into()))
                    .collect(),
            ),
        )
    }

    #[tokio::test]
    async fn verbalization_produces_readable_chain() {
        let llm = scripted(vec![
            "Ticket T-1 surfaces fact-alpha; Policy P-1 surfaces fact-beta.",
        ]);
        let pipeline = VerbalizationPipeline::new(llm);
        let (text, meta) = pipeline.verbalize(&path_with_two_hops()).await.unwrap();
        assert!(text.contains("fact-alpha"));
        assert!(text.contains("fact-beta"));
        assert!(meta.passed, "critic must accept on first draft");
        assert_eq!(meta.iterations_used, 0);
    }

    #[tokio::test]
    async fn refiner_revisits_when_first_draft_misses_facts() {
        let llm = scripted(vec![
            "Ticket T-1 surfaces fact-alpha but the second fact is unmentioned.",
            "Ticket T-1 surfaces fact-alpha and Policy P-1 surfaces fact-beta.",
        ]);
        let pipeline = VerbalizationPipeline::new(llm);
        let (text, meta) = pipeline.verbalize(&path_with_two_hops()).await.unwrap();
        assert!(text.contains("fact-alpha"));
        assert!(text.contains("fact-beta"));
        assert!(meta.passed);
        assert_eq!(meta.iterations_used, 1, "one refinement round needed");
    }

    #[tokio::test]
    async fn quality_loop_terminates_within_bounds() {
        // Two replies + cap iterations at 2 → initial draft + one
        // refine, then the second critique exhausts the budget
        // without another refine call. Script length matches the
        // exact call count.
        let llm = scripted(vec!["Mentions only fact-alpha.", "Still only fact-alpha."]);
        let pipeline = VerbalizationPipeline::new(llm).with_max_iterations(2);
        let (_text, meta) = pipeline.verbalize(&path_with_two_hops()).await.unwrap();
        assert!(!meta.passed);
        assert_eq!(meta.iterations_used, 2);
        assert_eq!(meta.max_iterations, 2);
    }

    #[tokio::test]
    async fn empty_path_passes_immediately() {
        let llm = scripted(vec!["(no facts)"]);
        let pipeline = VerbalizationPipeline::new(llm);
        let (_text, meta) = pipeline
            .verbalize(&RetrievalPath::with_hops(Vec::new()))
            .await
            .unwrap();
        assert!(meta.passed);
        assert_eq!(meta.iterations_used, 0);
    }

    #[tokio::test]
    async fn zero_max_iterations_rejected() {
        let llm = scripted(vec!["draft"]);
        let pipeline = VerbalizationPipeline::new(llm).with_max_iterations(0);
        let err = pipeline.verbalize(&path_with_two_hops()).await.unwrap_err();
        assert!(matches!(err, SpecError::Config(_)));
    }

    #[test]
    fn render_hops_lists_each_fact_id() {
        let path = path_with_two_hops();
        let rendered = render_hops(&path);
        assert!(rendered.contains("fact-alpha"));
        assert!(rendered.contains("fact-beta"));
        assert!(
            !rendered.contains("provenance=signed"),
            "unsigned hops must not carry the provenance marker"
        );
    }

    #[tokio::test]
    async fn critic_accepts_when_all_facts_present() {
        let path = path_with_two_hops();
        let critic = VerbalizationCritic::new(&path);
        let crit = critic
            .evaluate(&"fact-alpha and fact-beta are both here".to_string(), 0)
            .await
            .unwrap();
        assert!(crit.pass);
    }

    #[tokio::test]
    async fn critic_lists_missing_facts_in_feedback() {
        let path = path_with_two_hops();
        let critic = VerbalizationCritic::new(&path);
        let crit = critic
            .evaluate(&"only fact-alpha is mentioned".to_string(), 0)
            .await
            .unwrap();
        assert!(!crit.pass);
        assert!(
            crit.feedback.contains("fact-beta"),
            "feedback must name the missing fact: {}",
            crit.feedback
        );
    }
}