ai_tokenopt 0.5.9

//! Progressive tool schema compression.
//!
//! Tracks which tools have been seen in a conversation and progressively
//! strips detail on subsequent appearances: full definition on first use,
//! minimal (name + types only) on repeat.

use std::collections::HashSet;

use crate::types::{ParameterProperty, ToolDefinition, ToolParameters};

/// Tracks which tools have been sent to the LLM in a conversation.
///
/// Used to progressively reduce tool schema verbosity across turns.
#[derive(Debug, Clone, Default)]
pub struct ToolUsageTracker {
    /// Names of tools that have been included in at least one request
    seen_tools: HashSet<String>,
}

impl ToolUsageTracker {
    /// Create a new tracker with no history.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Mark tools as seen (should be called after each request).
    pub fn mark_seen(&mut self, tools: &[ToolDefinition]) {
        for tool in tools {
            self.seen_tools.insert(tool.name.clone());
        }
    }

    /// Check if a tool has been seen before.
    #[must_use]
    pub fn is_seen(&self, name: &str) -> bool {
        self.seen_tools.contains(name)
    }

    /// Reset tracker (e.g. on new conversation).
    pub fn reset(&mut self) {
        self.seen_tools.clear();
    }
}

/// Compress tool definitions progressively based on usage history.
///
/// - **First appearance**: full definition (unchanged)
/// - **Subsequent appearances**: name + parameter names + types only
///   (descriptions stripped for all parameters and the tool itself)
///
/// This runs *after* existing `select_tools()` and `compress_tool_definitions()`.
#[must_use]
pub fn compress_progressively(
    tools: &[ToolDefinition],
    tracker: &ToolUsageTracker,
) -> Vec<ToolDefinition> {
    tools
        .iter()
        .map(|tool| {
            if tracker.is_seen(&tool.name) {
                // Minimal: strip all descriptions
                ToolDefinition {
                    name: tool.name.clone(),
                    description: String::new(),
                    parameters: ToolParameters {
                        schema_type: tool.parameters.schema_type.clone(),
                        properties: tool
                            .parameters
                            .properties
                            .iter()
                            .map(|(k, v)| {
                                (
                                    k.clone(),
                                    ParameterProperty {
                                        param_type: v.param_type.clone(),
                                        description: String::new(),
                                        enum_values: v.enum_values.clone(),
                                    },
                                )
                            })
                            .collect(),
                        required: tool.parameters.required.clone(),
                    },
                    icon: tool.icon.clone(),
                }
            } else {
                tool.clone()
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;

    fn make_tool(name: &str) -> ToolDefinition {
        let mut props = HashMap::new();
        props.insert(
            "query".to_string(),
            ParameterProperty {
                param_type: "string".to_string(),
                description: "The search query to execute".to_string(),
                enum_values: Vec::new(),
            },
        );
        ToolDefinition {
            name: name.to_string(),
            description: format!("A tool named {name} that does something useful"),
            parameters: ToolParameters {
                schema_type: "object".to_string(),
                properties: props,
                required: vec!["query".to_string()],
            },
            icon: None,
        }
    }

    #[test]
    fn first_appearance_is_unchanged() {
        let tracker = ToolUsageTracker::new();
        let tools = vec![make_tool("search")];
        let result = compress_progressively(&tools, &tracker);
        assert_eq!(result[0].description, tools[0].description);
        assert_eq!(
            result[0].parameters.properties["query"].description,
            "The search query to execute"
        );
    }

    #[test]
    fn second_appearance_stripped() {
        let mut tracker = ToolUsageTracker::new();
        let tools = vec![make_tool("search")];
        tracker.mark_seen(&tools);

        let result = compress_progressively(&tools, &tracker);
        assert!(result[0].description.is_empty());
        assert!(
            result[0].parameters.properties["query"]
                .description
                .is_empty()
        );
        // Type is preserved
        assert_eq!(
            result[0].parameters.properties["query"].param_type,
            "string"
        );
        assert_eq!(result[0].name, "search");
    }

    #[test]
    fn mixed_seen_and_unseen() {
        let mut tracker = ToolUsageTracker::new();
        let tools = vec![make_tool("search"), make_tool("weather")];
        tracker.mark_seen(&[make_tool("search")]);

        let result = compress_progressively(&tools, &tracker);
        // search: stripped
        assert!(result[0].description.is_empty());
        // weather: full
        assert!(!result[1].description.is_empty());
    }

    #[test]
    fn new_tool_mid_conversation() {
        let mut tracker = ToolUsageTracker::new();
        tracker.mark_seen(&[make_tool("search")]);

        // New tool appears
        let tools = vec![make_tool("search"), make_tool("calendar")];
        let result = compress_progressively(&tools, &tracker);
        assert!(result[0].description.is_empty()); // seen
        assert!(!result[1].description.is_empty()); // new
    }

    #[test]
    fn reset_clears_history() {
        let mut tracker = ToolUsageTracker::new();
        tracker.mark_seen(&[make_tool("search")]);
        assert!(tracker.is_seen("search"));

        tracker.reset();
        assert!(!tracker.is_seen("search"));
    }

    #[test]
    fn enum_values_preserved_in_minimal() {
        let mut props = HashMap::new();
        props.insert(
            "format".to_string(),
            ParameterProperty {
                param_type: "string".to_string(),
                description: "Output format to use".to_string(),
                enum_values: vec!["json".to_string(), "text".to_string()],
            },
        );
        let tool = ToolDefinition {
            name: "convert".to_string(),
            description: "Conversion tool".to_string(),
            parameters: ToolParameters {
                schema_type: "object".to_string(),
                properties: props,
                required: vec!["format".to_string()],
            },
            icon: None,
        };

        let mut tracker = ToolUsageTracker::new();
        tracker.mark_seen(std::slice::from_ref(&tool));

        let result = compress_progressively(&[tool], &tracker);
        assert_eq!(
            result[0].parameters.properties["format"].enum_values,
            vec!["json", "text"]
        );
    }
}