ai_tokenopt 0.5.6

//! Tool call chain collapsing.
//!
//! When a sequence of tool calls resolves to a final answer, the
//! intermediate tool outputs often become redundant. This module
//! collapses consecutive tool messages into a summary, keeping only
//! the final relevant result.

use crate::estimator::TokenEstimator;
use crate::types::{ChatMessage, MessageRole};

/// Create a tool message (abstracts standalone vs. pisovereign constructors).
fn make_tool_message(content: impl Into<String>) -> ChatMessage {
    #[cfg(not(feature = "pisovereign"))]
    {
        ChatMessage::tool(content)
    }
    #[cfg(feature = "pisovereign")]
    {
        ChatMessage::tool("collapsed", content)
    }
}

/// Result of collapsing tool call chains.
#[derive(Debug)]
pub struct ChainCollapseResult {
    /// Messages after collapsing tool chains
    pub messages: Vec<ChatMessage>,
    /// Number of tool messages collapsed
    pub collapsed_count: usize,
    /// Estimated tokens saved by collapsing
    pub tokens_saved: u32,
}

/// Maximum length (characters) for a tool result summary line.
const SUMMARY_MAX_CHARS: usize = 120;

/// Collapse consecutive tool messages into a single summary message.
///
/// A "tool chain" is defined as 2+ consecutive [`MessageRole::Tool`]
/// messages. These are collapsed into a single tool message containing
/// a one-line summary of each collapsed result.
///
/// If a tool result is short enough (≤ `SUMMARY_MAX_CHARS`), it is
/// kept verbatim. Longer results are truncated with `...`.
///
/// # Examples
///
/// ```ignore
/// use ai_tokenopt::tools::chain_collapser::{collapse_tool_chains, ChainCollapseResult};
/// use ai_tokenopt::types::{ChatMessage, MessageRole};
///
/// let messages = vec![
///     ChatMessage::user("Check the weather"),
///     ChatMessage::tool("API call: get_weather(Berlin)"),
///     ChatMessage::tool("Result: {\"temp\": 22, \"condition\": \"sunny\"}"),
///     ChatMessage::tool("Processed: Temperature is 22°C, sunny"),
///     ChatMessage::assistant("It's 22°C and sunny in Berlin!"),
/// ];
/// let result = collapse_tool_chains(&messages);
/// assert!(result.collapsed_count > 0);
/// ```
#[must_use]
pub fn collapse_tool_chains(messages: &[ChatMessage]) -> ChainCollapseResult {
    if messages.len() < 2 {
        return ChainCollapseResult {
            messages: messages.to_vec(),
            collapsed_count: 0,
            tokens_saved: 0,
        };
    }

    let mut result: Vec<ChatMessage> = Vec::with_capacity(messages.len());
    let mut collapsed_count: usize = 0;
    let mut tokens_saved: u32 = 0;
    let mut i = 0;

    while i < messages.len() {
        // Check if this starts a tool chain (2+ consecutive Tool messages)
        if messages[i].role == MessageRole::Tool {
            let chain_start = i;
            while i < messages.len() && messages[i].role == MessageRole::Tool {
                i += 1;
            }
            let chain_end = i;
            let chain_len = chain_end - chain_start;

            if chain_len >= 2 {
                // Collapse the chain
                let original_tokens: u32 = messages[chain_start..chain_end]
                    .iter()
                    .map(|m| TokenEstimator::estimate_tokens(&m.content))
                    .sum();

                let summary = summarize_tool_chain(&messages[chain_start..chain_end]);
                let summary_tokens = TokenEstimator::estimate_tokens(&summary);

                result.push(make_tool_message(summary));
                collapsed_count += chain_len - 1;
                tokens_saved += original_tokens.saturating_sub(summary_tokens);
            } else {
                // Single tool message — keep as-is
                result.push(messages[chain_start].clone());
            }
        } else {
            result.push(messages[i].clone());
            i += 1;
        }
    }

    ChainCollapseResult {
        messages: result,
        collapsed_count,
        tokens_saved,
    }
}

/// Create a one-line-per-entry summary of a tool chain.
fn summarize_tool_chain(tool_messages: &[ChatMessage]) -> String {
    let mut lines: Vec<String> = Vec::with_capacity(tool_messages.len());

    for (idx, msg) in tool_messages.iter().enumerate() {
        let truncated = if msg.content.len() <= SUMMARY_MAX_CHARS {
            msg.content.clone()
        } else {
            format!(
                "{}...",
                &msg.content[..SUMMARY_MAX_CHARS.min(msg.content.len())]
            )
        };
        lines.push(format!("[step {}] {}", idx + 1, truncated));
    }

    lines.join("\n")
}

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

    #[test]
    fn no_tool_messages_unchanged() {
        let msgs = vec![ChatMessage::user("Hello"), ChatMessage::assistant("Hi!")];
        let result = collapse_tool_chains(&msgs);
        assert_eq!(result.messages.len(), 2);
        assert_eq!(result.collapsed_count, 0);
        assert_eq!(result.tokens_saved, 0);
    }

    #[test]
    fn single_tool_message_unchanged() {
        let msgs = vec![
            ChatMessage::user("Check weather"),
            make_tool_message("Sunny, 22°C"),
            ChatMessage::assistant("It's sunny!"),
        ];
        let result = collapse_tool_chains(&msgs);
        assert_eq!(result.messages.len(), 3);
        assert_eq!(result.collapsed_count, 0);
    }

    #[test]
    fn two_tool_messages_collapsed() {
        let msgs = vec![
            ChatMessage::user("Check weather"),
            make_tool_message("API call: get_weather"),
            make_tool_message("Result: sunny 22°C"),
            ChatMessage::assistant("It's sunny!"),
        ];
        let result = collapse_tool_chains(&msgs);
        assert_eq!(result.messages.len(), 3); // user, collapsed_tool, assistant
        assert_eq!(result.collapsed_count, 1);
        assert!(result.messages[1].content.contains("[step 1]"));
        assert!(result.messages[1].content.contains("[step 2]"));
    }

    #[test]
    fn three_tool_messages_collapsed() {
        let msgs = vec![
            make_tool_message("step one data"),
            make_tool_message("step two data"),
            make_tool_message("step three data"),
        ];
        let result = collapse_tool_chains(&msgs);
        assert_eq!(result.messages.len(), 1);
        assert_eq!(result.collapsed_count, 2);
    }

    #[test]
    fn multiple_separate_chains() {
        let msgs = vec![
            make_tool_message("chain1 a"),
            make_tool_message("chain1 b"),
            ChatMessage::user("middle"),
            make_tool_message("chain2 a"),
            make_tool_message("chain2 b"),
        ];
        let result = collapse_tool_chains(&msgs);
        // 2 collapsed chains + user = 3
        assert_eq!(result.messages.len(), 3);
        assert_eq!(result.collapsed_count, 2);
    }

    #[test]
    fn long_tool_output_truncated() {
        let long_content = "x".repeat(200);
        let msgs = vec![make_tool_message(&long_content), make_tool_message("short")];
        let result = collapse_tool_chains(&msgs);
        let summary = &result.messages[0].content;
        assert!(summary.contains("..."));
    }

    #[test]
    fn empty_messages() {
        let msgs: Vec<ChatMessage> = Vec::new();
        let result = collapse_tool_chains(&msgs);
        assert!(result.messages.is_empty());
    }

    #[test]
    fn tokens_saved_positive_for_long_chains() {
        let msgs = vec![
            make_tool_message("This is a fairly long tool output that takes up tokens"),
            make_tool_message("This is another fairly long tool output with different content"),
            make_tool_message("Yet another tool output with even more redundant data included"),
        ];
        let result = collapse_tool_chains(&msgs);
        // Summary should be shorter than original (compressed via one-line format)
        assert!(result.collapsed_count > 0);
    }
}