lc-agents 0.21.0

// lc-agents/src/crag/tests.rs
//! Unit tests for `CorrectiveRAGAgent`.

use super::*;
use async_trait::async_trait;
use futures_util::Stream;
use lc_core::language_models::{BaseChatModel, BaseLanguageModel, LLMResult, StreamChunk};
use lc_core::runnables::{Runnable, RunnableConfig};
use lc_core::tools::ToolError;
use lc_rag::RetrieverError;
use lc_schema::Message;
use lc_vector_stores::{Document, SearchResult};
use std::pin::Pin;

/// Error type for mock chat model.
#[derive(Debug, thiserror::Error)]
#[error("mock error: {0}")]
struct MockError(String);

// === Mock LLM ===

/// A mock chat model that returns configurable responses in sequence.
#[derive(Debug, Clone)]
struct MockChatModel {
    responses: Vec<String>,
    call_count: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}

impl MockChatModel {
    fn new(responses: Vec<&str>) -> Self {
        Self {
            responses: responses.iter().map(|s| s.to_string()).collect(),
            call_count: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        }
    }
}

#[async_trait]
impl Runnable<Vec<Message>, LLMResult> for MockChatModel {
    type Error = MockError;

    async fn invoke(
        &self,
        _input: Vec<Message>,
        _config: Option<RunnableConfig>,
    ) -> Result<LLMResult, Self::Error> {
        let idx = self
            .call_count
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let response = self
            .responses
            .get(idx)
            .unwrap_or(&"relevant".to_string())
            .clone();
        Ok(LLMResult {
            content: response,
            model: "mock".to_string(),
            token_usage: None,
            tool_calls: None,
            thinking_content: None,
        })
    }
}

#[async_trait]
impl BaseLanguageModel<Vec<Message>, LLMResult> for MockChatModel {
    fn model_name(&self) -> &str {
        "mock"
    }
    fn get_num_tokens(&self, text: &str) -> usize {
        text.split_whitespace().count()
    }
    fn with_temperature(self, _temp: f32) -> Self {
        self
    }
    fn with_max_tokens(self, _max: usize) -> Self {
        self
    }
}

#[async_trait]
impl BaseChatModel for MockChatModel {
    async fn chat(
        &self,
        _messages: Vec<Message>,
        _config: Option<RunnableConfig>,
    ) -> Result<LLMResult, Self::Error> {
        let idx = self
            .call_count
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let response = self
            .responses
            .get(idx)
            .unwrap_or(&"relevant".to_string())
            .clone();
        Ok(LLMResult {
            content: response,
            model: "mock".to_string(),
            token_usage: None,
            tool_calls: None,
            thinking_content: None,
        })
    }

    async fn stream_chat(
        &self,
        _messages: Vec<Message>,
        _config: Option<RunnableConfig>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
    {
        Err(MockError("streaming not supported".to_string()))
    }
}

// === Mock Retriever ===

#[derive(Debug, Clone)]
struct MockRetriever {
    documents: Vec<Document>,
}

impl MockRetriever {
    fn new(documents: Vec<Document>) -> Self {
        Self { documents }
    }
}

#[async_trait]
impl RetrieverTrait for MockRetriever {
    async fn retrieve(&self, _query: &str, k: usize) -> Result<Vec<Document>, RetrieverError> {
        Ok(self.documents.iter().take(k).cloned().collect())
    }

    async fn retrieve_with_scores(
        &self,
        query: &str,
        k: usize,
    ) -> Result<Vec<SearchResult>, RetrieverError> {
        let docs = self.retrieve(query, k).await?;
        Ok(docs
            .into_iter()
            .enumerate()
            .map(|(i, doc)| SearchResult {
                document: doc,
                score: 1.0 - (i as f32 * 0.1),
            })
            .collect())
    }

    async fn add_documents(&self, _documents: Vec<Document>) -> Result<(), RetrieverError> {
        Ok(())
    }
}

// === Mock Web Tool ===

struct MockWebTool;

#[async_trait]
impl BaseTool for MockWebTool {
    fn name(&self) -> &str {
        "web_search"
    }
    fn description(&self) -> &str {
        "Search the web"
    }
    async fn run(&self, _input: String) -> Result<String, ToolError> {
        Ok("Web search result: CRAG is Corrective RAG.".to_string())
    }
}

// === Tests ===

#[tokio::test]
async fn test_crag_agent_high_score_documents() {
    let llm = MockChatModel::new(vec![
        "Relevance: relevant\nScore: 0.9\nReasoning: Directly addresses the query.",
        "Relevance: relevant\nScore: 0.8\nReasoning: Closely related.",
        "Rust is a systems programming language focused on safety and performance.",
        "grounded",
    ]);

    let retriever = MockRetriever::new(vec![
        Document::new("Rust is a systems programming language."),
        Document::new("Rust emphasizes memory safety."),
    ]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grade_threshold(0.5)
        .with_hallucination_check(true);

    let result = agent.invoke("What is Rust?").await.unwrap();
    assert!(!result.answer.is_empty());
    assert!(result.grounded);
    assert_eq!(result.sources.len(), 2);
    assert_eq!(result.grade_scores.len(), 2);
    assert!(result.grade_scores[0] >= 0.5);
}

#[tokio::test]
async fn test_crag_agent_low_score_triggers_correction() {
    // generate_alternatives returns 3 alternative queries
    let llm = MockChatModel::new(vec![
        "Relevance: irrelevant\nScore: 0.1\nReasoning: Not related.",
        "Relevance: irrelevant\nScore: 0.2\nReasoning: Barely related.",
        "1. What are the key features of Rust?\n2. Rust programming language overview\n3. Rust memory safety and performance",
        "Relevance: relevant\nScore: 0.9\nReasoning: Directly addresses.",
        "Relevance: relevant\nScore: 0.8\nReasoning: Closely related.",
        "Rust provides memory safety without garbage collection.",
        "grounded",
    ]);

    let retriever = MockRetriever::new(vec![
        Document::new("Rust provides memory safety guarantees."),
        Document::new("Rust has zero-cost abstractions."),
    ]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grade_threshold(0.5)
        .with_hallucination_check(true);

    let result = agent.invoke("Tell me about Rust").await.unwrap();
    assert!(!result.answer.is_empty());
    assert!(result.grounded);
}

#[tokio::test]
async fn test_crag_agent_with_web_fallback() {
    let llm = MockChatModel::new(vec![
        "Relevance: irrelevant\nScore: 0.1\nReasoning: Not related.",
        "1. What is CRAG in AI?\n2. Corrective RAG technique\n3. CRAG methodology overview",
        "Relevance: relevant\nScore: 0.9\nReasoning: Direct match.",
        "CRAG stands for Corrective RAG.",
        "grounded",
    ]);

    let retriever = MockRetriever::new(vec![Document::new(
        "CRAG is a retrieval-augmented generation technique.",
    )]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grade_threshold(0.5)
        .with_web_fallback(Box::new(MockWebTool))
        .with_hallucination_check(true);

    let result = agent.invoke("What is CRAG?").await.unwrap();
    assert!(!result.answer.is_empty());
}

#[tokio::test]
async fn test_crag_agent_no_documents_retrieved() {
    let llm = MockChatModel::new(vec![]);
    let retriever = MockRetriever::new(vec![]);

    let agent = CorrectiveRAGAgent::new(llm, retriever);

    let result = agent.invoke("What is Rust?").await;
    assert!(result.is_err());
    match result.unwrap_err() {
        CRAGError::NoDocumentsRetrieved => {}
        other => panic!("Expected NoDocumentsRetrieved, got: {}", other),
    }
}

#[tokio::test]
async fn test_crag_agent_hallucination_detected() {
    let llm = MockChatModel::new(vec![
        "Relevance: relevant\nScore: 0.9\nReasoning: Direct match.",
        "Rust was invented by aliens in 3020.",
        "not grounded",
    ]);

    let retriever = MockRetriever::new(vec![Document::new(
        "Rust was created by Graydon Hoare in 2010.",
    )]);

    let agent = CorrectiveRAGAgent::new(llm, retriever).with_hallucination_check(true);

    let result = agent.invoke("Who created Rust?").await.unwrap();
    assert!(!result.grounded);
}

#[tokio::test]
async fn test_crag_agent_hallucination_check_disabled() {
    let llm = MockChatModel::new(vec![
        "Relevance: relevant\nScore: 0.9\nReasoning: Direct match.",
        "Rust is great.",
    ]);

    let retriever = MockRetriever::new(vec![Document::new("Rust is a programming language.")]);

    let agent = CorrectiveRAGAgent::new(llm, retriever).with_hallucination_check(false);

    let result = agent.invoke("What is Rust?").await.unwrap();
    // grounded defaults to true when check is disabled
    assert!(result.grounded);
}

#[test]
fn test_crag_result_fields() {
    let result = CRAGResult {
        answer: "Test answer".to_string(),
        grounded: true,
        sources: vec![Document::new("Source 1")],
        grade_scores: vec![0.9],
        grade_reasoning: vec![Some("Directly relevant".to_string())],
    };
    assert_eq!(result.answer, "Test answer");
    assert!(result.grounded);
    assert_eq!(result.sources.len(), 1);
    assert_eq!(result.grade_scores.len(), 1);
    assert_eq!(result.grade_reasoning.len(), 1);
}

#[test]
fn test_crag_error_display() {
    let err = CRAGError::NoDocumentsRetrieved;
    assert!(err.to_string().contains("No documents retrieved"));

    let err = CRAGError::GenerationError("timeout".to_string());
    assert!(err.to_string().contains("timeout"));
}

#[test]
fn test_grade_threshold_clamping() {
    let llm = MockChatModel::new(vec![]);
    let retriever = MockRetriever::new(vec![Document::new("test")]);

    let agent = CorrectiveRAGAgent::new(llm, retriever).with_grade_threshold(1.5);
    assert!((agent.grade_threshold - 1.0).abs() < f64::EPSILON);

    let agent = agent.with_grade_threshold(-0.5);
    assert!((agent.grade_threshold - 0.0).abs() < f64::EPSILON);
}

#[test]
fn test_default_grade_threshold_is_0_6() {
    let llm = MockChatModel::new(vec![]);
    let retriever = MockRetriever::new(vec![Document::new("test")]);

    let agent = CorrectiveRAGAgent::new(llm, retriever);
    assert!((agent.grade_threshold - 0.6).abs() < f64::EPSILON);
}

#[test]
fn test_grader_llm_is_stored_when_set() {
    let llm = MockChatModel::new(vec![]);
    let retriever = MockRetriever::new(vec![Document::new("test")]);

    let grader = MockChatModel::new(vec!["grounded"]);
    let agent = CorrectiveRAGAgent::new(llm, retriever).with_grader_llm(grader);
    assert!(agent.grader_llm.is_some());
}

#[tokio::test]
async fn test_crag_agent_with_grader_llm() {
    // Grader LLM returns "not grounded" -> answer marked as ungrounded
    let llm = MockChatModel::new(vec![
        "Relevance: relevant\nScore: 0.9\nReasoning: Direct match.",
        "Rust was invented by aliens.",
        "not grounded",
    ]);
    let grader = MockChatModel::new(vec!["not grounded"]);

    let retriever = MockRetriever::new(vec![Document::new("Rust was created by Graydon Hoare.")]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grader_llm(grader)
        .with_hallucination_check(true);

    let result = agent.invoke("Who created Rust?").await.unwrap();
    assert!(!result.grounded);
}

#[tokio::test]
async fn test_crag_agent_stream_high_score() {
    use futures_util::StreamExt;

    let llm = MockChatModel::new(vec![
        "Relevance: relevant\nScore: 0.9\nReasoning: Directly addresses the query.",
        "Relevance: relevant\nScore: 0.8\nReasoning: Closely related.",
        "Rust is a systems programming language.",
        "grounded",
    ]);

    let retriever = MockRetriever::new(vec![
        Document::new("Rust is a systems programming language."),
        Document::new("Rust emphasizes memory safety."),
    ]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grade_threshold(0.5)
        .with_hallucination_check(true);

    let stream = agent.stream("What is Rust?").await.unwrap();
    let events: Vec<_> = stream.collect().await;

    // Should have: retrieving, retrieved, grading, graded, generating, hallucination_check, FinalAnswer
    assert!(
        events.len() >= 5,
        "Expected at least 5 events, got {}",
        events.len()
    );

    // First event should be retrieving
    assert!(matches!(
        &events[0],
        crate::streaming::AgentStreamEvent::PipelineStep { step, .. } if step == "retrieving"
    ));

    // Last event should be FinalAnswer
    assert!(matches!(
        events.last().unwrap(),
        crate::streaming::AgentStreamEvent::FinalAnswer { .. }
    ));
}

#[tokio::test]
async fn test_crag_agent_stream_low_score_correction() {
    use futures_util::StreamExt;

    let llm = MockChatModel::new(vec![
        "Relevance: irrelevant\nScore: 0.1\nReasoning: Not related.",
        "1. Rust features\n2. Rust language\n3. Rust memory safety",
        "Relevance: relevant\nScore: 0.9\nReasoning: Directly addresses.",
        "Rust provides memory safety.",
        "grounded",
    ]);

    let retriever = MockRetriever::new(vec![Document::new(
        "Rust provides memory safety guarantees.",
    )]);

    let agent = CorrectiveRAGAgent::new(llm, retriever)
        .with_grade_threshold(0.5)
        .with_hallucination_check(false);

    let stream = agent.stream("Tell me about Rust").await.unwrap();
    let events: Vec<_> = stream.collect().await;

    // Should contain correcting + corrected events
    let step_names: Vec<&str> = events
        .iter()
        .filter_map(|e| match e {
            crate::streaming::AgentStreamEvent::PipelineStep { step, .. } => Some(step.as_str()),
            _ => None,
        })
        .collect();

    assert!(
        step_names.contains(&"correcting"),
        "Expected 'correcting' step, got: {:?}",
        step_names
    );
    assert!(
        step_names.contains(&"corrected"),
        "Expected 'corrected' step, got: {:?}",
        step_names
    );
}