monothink 0.2.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 std::collections::HashMap;

use crate::core::conversation::{ConversationItem, ConversationItemKind};
use unicode_width::UnicodeWidthStr;

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: &str, history: &[ConversationItem]) -> Vec<ConversationItem>;
}

#[derive(Debug, PartialEq, Clone)]
pub struct SlidingWindowCompactor {
    // the approximate number of characters in a token
    chars_per_token: f64,
    // the maximum limit (in tokens) in the context window before compaction is needed
    context_limit: u64,
}

impl SlidingWindowCompactor {
    pub fn new(chars_per_token: f64, context_limit: u64) -> Self {
        Self {
            chars_per_token,
            context_limit,
        }
    }
}

impl Compactor for SlidingWindowCompactor {
    fn compact(&self, prompt: &str, history: &[ConversationItem]) -> Vec<ConversationItem> {
        // @todo: this is to lossy. find a better alternative (e.g. a fraction)
        let mut character_limit = (self.chars_per_token * (self.context_limit as f64)) as u64;
        character_limit -= prompt.width() as u64;

        let mut truncated_history: Vec<ConversationItem> = history
            .iter()
            .rev()
            .filter(|item| !is_empty_agent_message(item))
            .take_while(|item| {
                if character_limit >= item.context_text_length() {
                    character_limit -= item.context_text_length();
                    true
                } else {
                    false
                }
            })
            .map(|item| item.to_owned())
            .collect();
        truncated_history.reverse();

        remove_dangling_tool_responses(&mut truncated_history)
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct CutOffCompactor {
    // the approximate number of characters in a token
    chars_per_token: f64,
    // the maximum limit (in tokens) in the context window before compaction is needed
    context_limit: u64,
    // the size (in tokens) of history to keep verbatim
    preservation_size: u64,
}

impl CutOffCompactor {
    pub fn new(chars_per_token: f64, context_limit: u64, preservation_size: u64) -> Self {
        Self {
            chars_per_token,
            context_limit,
            preservation_size,
        }
    }
}

impl Compactor for CutOffCompactor {
    fn compact(&self, prompt: &str, history: &[ConversationItem]) -> Vec<ConversationItem> {
        // @todo: this is to lossy. find a better alternative (e.g. a fraction)
        let mut character_limit = (self.chars_per_token * (self.context_limit as f64)) as u64;
        character_limit -= prompt.width() as u64;

        let mut preserved_size = 0;
        let mut preserved_count = 0;
        let mut needs_cutoff = false;

        let relevant_history = history
            .iter()
            .rev()
            .filter(|item| !is_empty_agent_message(item));

        for item in relevant_history {
            if character_limit >= item.context_text_length() {
                character_limit -= item.context_text_length();
            } else {
                needs_cutoff = true;
                break;
            }

            if preserved_size + item.context_text_length() <= self.preservation_size {
                preserved_count += 1;
                preserved_size += item.context_text_length();
            }
        }

        if needs_cutoff {
            remove_dangling_tool_responses(&mut history[history.len() - preserved_count..].to_vec())
        } else {
            history.to_vec()
        }
    }
}

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

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

impl Compactor for ObservationMaskingCompactor {
    fn compact(&self, _prompt: &str, history: &[ConversationItem]) -> Vec<ConversationItem> {
        let mut verbatim_count = self.preservation_count;
        let mut observation_count = 1;
        let mut truncated_history: Vec<ConversationItem> = history
            .iter()
            .rev()
            .filter(|item| !is_empty_agent_message(item))
            .map(|item| match &item.content {
                ConversationItemKind::ToolOutput(tool_output_content) => {
                    if verbatim_count > 0 {
                        verbatim_count -= 1;
                        item.to_owned()
                    } else {
                        let mut new_content = tool_output_content.clone();
                        new_content.observation = Some((
                            observation_count,
                            format!("[Observation:{observation_count} - masked]"),
                        ));
                        observation_count += 1;

                        ConversationItem::new(ConversationItemKind::ToolOutput(new_content))
                    }
                }
                _ => item.to_owned(),
            })
            .collect();
        truncated_history.reverse();

        truncated_history
    }
}

fn is_empty_agent_message(item: &ConversationItem) -> bool {
    match &item.content {
        ConversationItemKind::AssistantMessage(message_content)
        | ConversationItemKind::AssistantReasoning(message_content) => {
            message_content.raw().trim().is_empty()
        }
        _ => false,
    }
}

// this assumes that history is remove from the end and calls always come before a response
fn remove_dangling_tool_responses(history: &mut Vec<ConversationItem>) -> Vec<ConversationItem> {
    let mut invocations = HashMap::new();
    let mut responses = HashMap::new();
    for item in history.iter().enumerate() {
        match &item.1.content {
            ConversationItemKind::ToolInvocation(tool_invocation_content) => {
                invocations.insert(tool_invocation_content.call_id.clone(), item.0);
            }
            ConversationItemKind::ToolOutput(tool_output_content) => {
                responses.insert(tool_output_content.call_id.clone(), item.0);
            }
            _ => {}
        }
    }

    for response in responses {
        if !invocations.contains_key(&response.0) {
            history.remove(response.1);
        }
    }
    todo!()
}