monothink 0.4.2

/*
* This file is part of monothink.
*
* monothink is free software: you can redistribute it and/or modify it under the terms
* of the GNU General Public License as published by the Free Software Foundation,
* either version 3 of the License, or (at your option) any later version.
*
* monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
* See the GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License along with Wavething.
* If not, see <https://www.gnu.org/licenses/>.

*/

use crate::llm::estimators::Estimator;
use rig::message::{AssistantContent, Message, Reasoning, ToolResultContent, UserContent};
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

pub trait Compactor {
    // returns the compacted history
    // @todo when implementing a summary or structured state compactor, move to async
    // then send request to llm and await response on a one_shot channel passed with the request
    // until then, avoid awaiting
    fn compact(&self, prompt: &Message, history: &[Message]) -> Vec<Message>;
}

#[derive(Debug, PartialEq, Clone)]
pub struct SlidingWindowCompactor<E: Estimator + Send + Sync> {
    // estimator used to calculate token length of a message
    estimator: E,
    // the maximum limit (in tokens) in the context window before compaction is needed
    context_limit: u64,
}

impl<E: Estimator + Send + Sync> SlidingWindowCompactor<E> {
    pub fn new(estimator: E, context_limit: u64) -> Self {
        Self {
            estimator,
            context_limit,
        }
    }
}

impl<E: Estimator + Send + Sync> Compactor for SlidingWindowCompactor<E> {
    fn compact(&self, prompt: &Message, history: &[Message]) -> Vec<Message> {
        let mut token_limit = self.context_limit;
        token_limit = token_limit.saturating_sub(self.estimator.estimate(prompt));
        token_limit = token_limit.saturating_sub(
            history
                .iter()
                .map(|msg| match msg {
                    Message::System { content: _ } => self.estimator.estimate(msg),
                    _ => 0,
                })
                .sum::<u64>(),
        );

        let mut truncated_history: Vec<Message> = history
            .iter()
            .rev()
            .filter(|msg| {
                let msg_len = self.estimator.estimate(msg);
                if token_limit >= msg_len && !matches!(msg, Message::System { content: _ }) {
                    token_limit -= msg_len;
                    true
                } else {
                    matches!(msg, Message::System { content: _ })
                }
            })
            .map(|item| item.to_owned())
            .collect();
        truncated_history.reverse();

        remove_dangling_tool_responses(&truncated_history)
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct ReasoningCapCompactor<E: Estimator + Send + Sync> {
    // estimator used to calculate token length of a message
    estimator: E,
    // the maximum limit (in tokens) in the context window before compaction is needed
    context_limit: u64,
    // the number of reasoning messages to keep verbatim
    preservation_count: u64,
}

impl<E: Estimator + Send + Sync> ReasoningCapCompactor<E> {
    pub fn new(estimator: E, context_limit: u64, preservation_count: u64) -> Self {
        Self {
            estimator,
            context_limit,
            preservation_count,
        }
    }
}

impl<E: Estimator + Send + Sync> Compactor for ReasoningCapCompactor<E> {
    fn compact(&self, prompt: &Message, history: &[Message]) -> Vec<Message> {
        let context_pressure = self.estimator.context_pressure(prompt, history);

        if context_pressure > self.context_limit {
            let mut max_ommitted_messages: u64 = 0;
            for msg in history {
                if let Message::Assistant { id: _, content } = msg {
                    for elem in content {
                        if let AssistantContent::Reasoning(_reasoning) = elem {
                            max_ommitted_messages += 1;
                        }
                    }
                }
            }
            max_ommitted_messages = max_ommitted_messages.saturating_sub(self.preservation_count);

            history
                .iter()
                .map(|msg| {
                    if let Message::Assistant { id, content } = msg {
                        Message::Assistant {
                            id: id.clone(),
                            content: content
                                .iter()
                                .map(|item| {
                                    if let AssistantContent::Reasoning(_reasoning) = item
                                        && max_ommitted_messages > 0
                                    {
                                        max_ommitted_messages -= 1;
                                        AssistantContent::Reasoning(Reasoning::summaries(vec![
                                            "Reasoning omitted due to age.".to_string(),
                                        ]))
                                    } else {
                                        item.to_owned()
                                    }
                                })
                                .collect(),
                        }
                    } else {
                        msg.to_owned()
                    }
                })
                .collect()
        } else {
            history.to_vec()
        }
    }
}

pub type SharedObservations = Arc<Mutex<HashMap<u64, Vec<ToolResultContent>>>>;

#[derive(Debug, Clone)]
pub struct ObservationMaskingCompactor {
    // the number of tool results to keep verbatim
    preservation_count: u64,
    observations: SharedObservations,
}

impl ObservationMaskingCompactor {
    pub fn new(preservation_count: u64, observations: SharedObservations) -> Self {
        Self {
            preservation_count,
            observations,
        }
    }
}

impl Compactor for ObservationMaskingCompactor {
    fn compact(&self, _prompt: &Message, history: &[Message]) -> Vec<Message> {
        let mut tool_result_count: u64 = 0;
        for msg in history {
            if let Message::User { content } = msg {
                for usr_content in content {
                    if let UserContent::ToolResult(_) = usr_content {
                        tool_result_count += 1;
                    }
                }
            }
        }
        let observation_limit: u64 = tool_result_count.saturating_sub(self.preservation_count);

        let mut observations = self.observations.lock().unwrap();
        let mut observation_count = 1;
        let masked_history: Vec<Message> = history
            .iter()
            .map(|item| match item {
                Message::User { content } => Message::User {
                    content: content
                        .iter()
                        .map(|user_item| match user_item {
                            rig::message::UserContent::ToolResult(tool_result) => {
                                if observation_count > observation_limit {
                                    user_item.to_owned()
                                } else {
                                    let mut new_content = tool_result.clone();
                                    new_content.content =
                                        vec![ToolResultContent::Text(rig::message::Text::new(
                                            format!("[Observation:{observation_count} - masked]"),
                                        ))];
                                    observations
                                        .insert(observation_count, tool_result.content.clone());
                                    observation_count += 1;
                                    rig::message::UserContent::ToolResult(new_content)
                                }
                            }
                            _ => user_item.clone(),
                        })
                        .collect(),
                },
                _ => item.to_owned(),
            })
            .collect();

        masked_history
    }
}

// this assumes that history removes from the end and calls always come before a response
fn remove_dangling_tool_responses(history: &[Message]) -> Vec<Message> {
    // get call id from all tool calls
    let mut calls = Vec::new();
    for item in history.iter() {
        if let Message::Assistant { id: _, content } = item {
            for elem in content {
                if let AssistantContent::ToolCall(tool_call) = elem {
                    calls.push(tool_call.id.clone())
                }
            }
        }
    }

    // go through all responses and drop all that are not present in calls
    history
        .iter()
        .filter_map(|msg| match msg {
            Message::User { content } => {
                let new_content: Vec<_> = content
                    .iter()
                    .filter(|elem| match elem {
                        UserContent::ToolResult(tool_result) => calls.contains(&tool_result.call),
                        _ => true,
                    })
                    .collect();
                if new_content.is_empty() {
                    None
                } else {
                    Some(msg.to_owned())
                }
            }
            _ => Some(msg.to_owned()),
        })
        .collect()
}