langchainrust 0.7.0

// src/memory/context_window/tests.rs
//! Tests for context_window module.

use super::*;
use std::sync::Arc;

use crate::core::language_models::{BaseChatModel, BaseLanguageModel, LLMResult};
use crate::core::runnables::{Runnable, RunnableConfig};
use crate::core::token_counter::TokenCounter;
use crate::language_models::openai::{OpenAIChat, OpenAIConfig};
use crate::schema::{Message, MessageType};
use async_trait::async_trait;
use futures_util::Stream;
use std::pin::Pin;
use tokio::sync::Mutex;

// ---- Mock TokenCounter for deterministic tests ----

/// A simple token counter that counts 1 token per character.
/// This makes it easy to reason about token budgets in tests.
#[derive(Debug)]
struct CharTokenCounter;

impl TokenCounter for CharTokenCounter {
    fn count_tokens(&self, text: &str) -> u32 {
        text.len() as u32
    }

    fn count_messages(&self, messages: &[Message]) -> u32 {
        let mut total = 0u32;
        for msg in messages {
            total += 4; // per-message overhead
            total += self.count_tokens(&msg.content);
        }
        total += 2; // conversation boundary
        total
    }
}

fn char_counter() -> Arc<dyn TokenCounter> {
    Arc::new(CharTokenCounter)
}

// ---- Mock LLM for Summarize strategy tests ----

#[derive(Debug)]
struct MockLLM {
    responses: Arc<Mutex<Vec<String>>>,
}

impl MockLLM {
    fn new(responses: Vec<String>) -> Self {
        Self {
            responses: Arc::new(Mutex::new(responses)),
        }
    }
}

impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
    fn model_name(&self) -> &str {
        "mock-llm"
    }

    fn get_num_tokens(&self, text: &str) -> usize {
        text.len()
    }

    fn with_temperature(self, _temp: f32) -> Self
    where
        Self: Sized,
    {
        self
    }

    fn with_max_tokens(self, _max: usize) -> Self
    where
        Self: Sized,
    {
        self
    }
}

#[async_trait]
impl Runnable<Vec<Message>, LLMResult> for MockLLM {
    type Error = std::convert::Infallible;

    async fn invoke(
        &self,
        _input: Vec<Message>,
        _config: Option<RunnableConfig>,
    ) -> Result<LLMResult, Self::Error> {
        let mut responses = self.responses.lock().await;
        let content = responses.pop().unwrap_or_else(|| "Summary".to_string());
        Ok(LLMResult {
            content,
            model: "mock-llm".to_string(),
            token_usage: None,
            tool_calls: None,
            thinking_content: None,
        })
    }
}

#[async_trait]
impl BaseChatModel for MockLLM {
    async fn chat(
        &self,
        messages: Vec<Message>,
        config: Option<RunnableConfig>,
    ) -> Result<LLMResult, Self::Error> {
        self.invoke(messages, config).await
    }

    async fn stream_chat(
        &self,
        _messages: Vec<Message>,
        _config: Option<RunnableConfig>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error> {
        unimplemented!("stream_chat not needed for tests")
    }
}

// ---- Helper to build messages ----

fn make_messages(contents: &[(&str, &str)]) -> Vec<Message> {
    contents
        .iter()
        .map(|(role, content)| match *role {
            "system" => Message::system(*content),
            "human" => Message::human(*content),
            "ai" => Message::ai(*content),
            _ => Message::human(*content),
        })
        .collect()
}

// ---- Tests ----

#[test]
fn test_new_creates_truncate_strategy() {
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(4096);
    assert_eq!(cw.max_tokens(), 4096);
}

#[test]
fn test_with_max_tokens() {
    let cw: ContextWindow<OpenAIChat> = ContextWindow::with_max_tokens(8192);
    assert_eq!(cw.max_tokens(), 8192);
}

#[tokio::test]
async fn test_fit_under_limit_returns_as_is() {
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(1000).with_counter(char_counter());

    let messages = make_messages(&[("human", "Hello"), ("ai", "Hi there")]);

    let result = cw.fit(messages).await.unwrap();
    assert_eq!(result.len(), 2);
}

#[tokio::test]
async fn test_fit_empty_messages() {
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(100).with_counter(char_counter());

    let result = cw.fit(vec![]).await.unwrap();
    assert!(result.is_empty());
}

#[tokio::test]
async fn test_truncate_preserves_system_messages() {
    // With CharTokenCounter: each message = 4 overhead + content length, + 2 boundary.
    // System: 4 + 7 = 11, Human1: 4 + 4 = 8, AI1: 4 + 4 = 8, Human2: 4 + 4 = 8, AI2: 4 + 4 = 8
    // Total = 11 + 8 + 8 + 8 + 8 + 2 = 45
    // Budget = 30: system(11) + AI2(8) + boundary(2) = 21, + Human2(8) = 29 <= 30
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(30).with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "You are"),
        ("human", "Q1?"),
        ("ai", "A1!"),
        ("human", "Q2?"),
        ("ai", "A2!"),
    ]);

    let result = cw.fit(messages).await.unwrap();

    // System message must be preserved.
    assert!(result
        .iter()
        .any(|m| matches!(m.message_type, MessageType::System)));
    // Most recent messages should be kept.
    assert!(result.iter().any(|m| m.content == "A2!"));
}

#[tokio::test]
async fn test_truncate_drops_oldest_first() {
    // Budget = 25: system(4+4=8) + AI(4+3=7) + boundary(2) = 17, + Human(4+3=7) = 24 <= 25
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(25).with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "Sys"),
        ("human", "Old question here"),
        ("ai", "Old answer here"),
        ("human", "New"),
        ("ai", "Ans"),
    ]);

    let result = cw.fit(messages).await.unwrap();

    // System message preserved.
    assert!(result.iter().any(|m| m.content == "Sys"));
    // Newest messages kept.
    assert!(result.iter().any(|m| m.content == "Ans"));
    // Old messages dropped.
    assert!(!result.iter().any(|m| m.content == "Old question here"));
}

#[tokio::test]
async fn test_truncate_only_system_messages() {
    // If only system messages exist and they fit, return them.
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(20).with_counter(char_counter());

    let messages = make_messages(&[("system", "Hello")]);
    // 4 + 5 + 2 = 11 <= 20

    let result = cw.fit(messages).await.unwrap();
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].content, "Hello");
}

#[tokio::test]
async fn test_truncate_system_only_over_budget() {
    // If system messages alone exceed the budget, truncate returns just system messages
    // (they are always preserved).
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(5).with_counter(char_counter());

    let messages = make_messages(&[("system", "Very long system prompt that exceeds budget")]);
    // 4 + 42 + 2 = 48 > 5

    let result = cw.fit(messages).await.unwrap();
    // System messages are always preserved even if over budget.
    assert_eq!(result.len(), 1);
}

#[tokio::test]
async fn test_summarize_replaces_old_messages() {
    // CharTokenCounter: per-message = 4 + content_len, + 2 boundary.
    // system("S"): 5, human("Q1"): 6, ai("A1"): 6, human("Q2"): 6, ai("A2"): 6,
    // human("Q3"): 6, ai("A3"): 6, human("Q4"): 6, ai("A4"): 6
    // Total = 5 + 8*6 + 2 = 55 > 40, so summarization is triggered.
    //
    // Budget 40: system(5) + placeholder(23) + ai("A4")(6) + boundary(2) = 36 <= 40.
    // keep_from_idx = 7 (keep ai("A4"), summarize Q1..Q4).
    // LLM returns "S." (2 chars). Summary = "[Conversation Summary] S." (24 chars) = 4+24=28 tokens.
    // Final: 5+28+6+2 = 41 > 40 => falls back to truncation.
    //
    // Use budget 50 instead:
    // system(5) + placeholder(23) + human("Q4")(6) + ai("A4")(6) + boundary(2) = 42 <= 50.
    // keep_from_idx = 6 (keep Q4, A4, summarize Q1..A3).
    // Summary = "[Conversation Summary] S." = 28 tokens. Final: 5+28+6+6+2 = 47 <= 50. Fits!
    let mock_llm = MockLLM::new(vec!["S.".to_string()]);

    let cw = ContextWindow::with_strategy(50, Strategy::summarize(mock_llm))
        .with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "S"),
        ("human", "Q1"),
        ("ai", "A1"),
        ("human", "Q2"),
        ("ai", "A2"),
        ("human", "Q3"),
        ("ai", "A3"),
        ("human", "Q4"),
        ("ai", "A4"),
    ]);

    let result = cw.fit(messages).await.unwrap();

    // System message preserved.
    assert!(result.iter().any(|m| m.content == "S"));

    // Should contain a summary message.
    let summary_msgs: Vec<&Message> = result
        .iter()
        .filter(|m| m.content.starts_with("[Conversation Summary]"))
        .collect();
    assert_eq!(summary_msgs.len(), 1);
    assert!(summary_msgs[0].content.contains("S"));
}

#[tokio::test]
async fn test_summarize_preserves_recent_messages() {
    // Budget = 50: system(5) + placeholder(23) + human("Q4")(6) + ai("A4")(6) + boundary(2) = 42 <= 50.
    // keep_from_idx = 6 (keep Q4, A4, summarize Q1..A3).
    // LLM returns "S." Summary = "[Conversation Summary] S." = 28 tokens.
    // Final: 5+28+6+6+2 = 47 <= 50. Fits!
    let mock_llm = MockLLM::new(vec!["S.".to_string()]);

    let cw = ContextWindow::with_strategy(50, Strategy::summarize(mock_llm))
        .with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "S"),
        ("human", "Q1"),
        ("ai", "A1"),
        ("human", "Q2"),
        ("ai", "A2"),
        ("human", "Q3"),
        ("ai", "A3"),
        ("human", "Q4"),
        ("ai", "A4"),
    ]);

    let result = cw.fit(messages).await.unwrap();

    // Recent messages should be preserved.
    assert!(result.iter().any(|m| m.content == "Q4"));
    assert!(result.iter().any(|m| m.content == "A4"));
}

#[tokio::test]
async fn test_summarize_with_custom_prompt() {
    // Budget = 50: system(5) + placeholder(23) + human("Q4")(6) + ai("A4")(6) + boundary(2) = 42 <= 50.
    // keep_from_idx = 6 (keep Q4, A4, summarize Q1..A3).
    // LLM returns "O." Summary = "[Conversation Summary] O." = 28 tokens.
    // Final: 5+28+6+6+2 = 47 <= 50. Fits!
    let mock_llm = MockLLM::new(vec!["O.".to_string()]);

    let cw = ContextWindow::with_strategy(
        50,
        Strategy::summarize_with_prompt(mock_llm, "Please compress: {conversation}\nCompressed:"),
    )
    .with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "S"),
        ("human", "Q1"),
        ("ai", "A1"),
        ("human", "Q2"),
        ("ai", "A2"),
        ("human", "Q3"),
        ("ai", "A3"),
        ("human", "Q4"),
        ("ai", "A4"),
    ]);

    let result = cw.fit(messages).await.unwrap();
    let summary_msgs: Vec<&Message> = result
        .iter()
        .filter(|m| m.content.starts_with("[Conversation Summary]"))
        .collect();
    assert_eq!(summary_msgs.len(), 1);
    assert!(summary_msgs[0].content.contains("O"));
}

#[tokio::test]
async fn test_summarize_no_non_system_messages() {
    let mock_llm = MockLLM::new(vec!["Should not be called".to_string()]);

    let cw: ContextWindow<MockLLM> =
        ContextWindow::with_strategy(50, Strategy::summarize(mock_llm))
            .with_counter(char_counter());

    let messages = make_messages(&[("system", "S")]);

    let result = cw.fit(messages).await.unwrap();
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].content, "S");
}

#[tokio::test]
async fn test_strategy_truncate_enum() {
    let cw = ContextWindow::with_strategy(100, Strategy::<OpenAIChat>::Truncate)
        .with_counter(char_counter());

    let messages = make_messages(&[("human", "Hello"), ("ai", "World")]);

    let result = cw.fit(messages).await.unwrap();
    assert_eq!(result.len(), 2);
}

#[test]
fn test_strategy_summarize_new() {
    let config = OpenAIConfig::default();
    let llm = OpenAIChat::new(config);
    let strategy: Strategy<OpenAIChat> = Strategy::summarize(llm);

    if let Strategy::Summarize { summary_prompt, .. } = &strategy {
        assert!(summary_prompt.contains("{conversation}"));
    } else {
        panic!("Expected Summarize variant");
    }
}

#[test]
fn test_strategy_summarize_with_custom_prompt() {
    let config = OpenAIConfig::default();
    let llm = OpenAIChat::new(config);
    let custom = "Custom: {conversation} ->";
    let strategy: Strategy<OpenAIChat> = Strategy::summarize_with_prompt(llm, custom);

    if let Strategy::Summarize { summary_prompt, .. } = &strategy {
        assert_eq!(summary_prompt, custom);
    } else {
        panic!("Expected Summarize variant");
    }
}

#[tokio::test]
async fn test_fit_with_real_tiktoken_counter() {
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(4096);

    let messages = make_messages(&[
        ("system", "You are a helpful assistant."),
        ("human", "Hello!"),
        ("ai", "Hi there! How can I help you?"),
    ]);

    // These short messages should easily fit within 4096 tokens.
    let result = cw.fit(messages).await.unwrap();
    assert_eq!(result.len(), 3);
}

#[tokio::test]
async fn test_truncate_preserves_order() {
    // Budget = 40: system(4+3=7) + human(4+3=7) + ai(4+3=7) + boundary(2) = 23 <= 40
    let cw: ContextWindow<OpenAIChat> = ContextWindow::new(40).with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "Sys"),
        ("human", "Old"),
        ("ai", "OldA"),
        ("human", "New"),
        ("ai", "NewA"),
    ]);

    let result = cw.fit(messages).await.unwrap();

    // Verify order: system first, then conversation in order.
    let types: Vec<&str> = result.iter().map(|m| m.type_str()).collect();
    // System should be first.
    assert_eq!(types[0], "system");
    // The rest should maintain human/ai alternation.
    for i in 1..types.len() {
        if i + 1 < types.len() {
            // Not strictly required, but for our test data this holds.
        }
    }
}

#[tokio::test]
async fn test_summarize_fallback_to_truncate() {
    // When the summary + recent messages still exceed the budget,
    // the method falls back to truncation.
    let mock_llm = MockLLM::new(vec![
        "A very long summary that will not fit in the small budget.".to_string(),
    ]);

    // Very small budget that even the summary won't fit.
    let cw: ContextWindow<MockLLM> =
        ContextWindow::with_strategy(20, Strategy::summarize(mock_llm))
            .with_counter(char_counter());

    let messages = make_messages(&[
        ("system", "S"),
        ("human", "Q1"),
        ("ai", "A1"),
        ("human", "Q2"),
        ("ai", "A2"),
    ]);

    let result = cw.fit(messages).await.unwrap();
    // Should still return some messages (truncation fallback).
    assert!(!result.is_empty());
}