cats 0.1.15

Coding Agent ToolS - A comprehensive toolkit for building AI-powered coding agents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use crate::core::{Tool, ToolArgs, ToolError, ToolResult};
use crate::tools::old::state::ToolState;
use anyhow::Result;
use std::fs;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

/// Tool to count tokens using tiktoken_rs cl100k_base encoding.
pub struct CountTokensTool {
    name: String,
}

impl CountTokensTool {
    pub fn new() -> Self {
        Self {
            name: "count_tokens".to_string(),
        }
    }

    /// Helper to count tokens from a string using cl100k_base tokenizer if enabled, otherwise fallback to whitespace count.
    #[cfg(feature = "tiktoken")]
    fn count_from_text(text: &str) -> usize {
        match tiktoken_rs::cl100k_base() {
            Ok(bpe) => {
                let tokens = bpe.encode_with_special_tokens(text);
                tokens.len()
            }
            Err(_) => text.split_whitespace().count(),
        }
    }

    #[cfg(not(feature = "tiktoken"))]
    fn count_from_text(text: &str) -> usize {
        // Fallback simple tokenization: split on whitespace
        text.split_whitespace().count()
    }
}

impl Tool for CountTokensTool {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        "Count tokens in a file or provided content using cl100k_base tokenizer"
    }

    fn signature(&self) -> &str {
        "count_tokens <file_path> OR count_tokens --content=<text>"
    }

    fn validate_args(&self, args: &ToolArgs) -> Result<(), ToolError> {
        if args.is_empty() {
            return Err(ToolError::InvalidArgs {
                message: "Usage: count_tokens <file_path> or count_tokens --content=<text>"
                    .to_string(),
            });
        }
        Ok(())
    }

    fn execute(&mut self, args: &ToolArgs, _state: &Arc<Mutex<ToolState>>) -> Result<ToolResult> {
        // If named arg `content` present, use it. Otherwise, take first positional arg as file path.
        if let Some(content) = args.get_named_arg("content") {
            let count = Self::count_from_text(content);
            return Ok(ToolResult::success_with_data(
                format!("Total tokens: {}", count),
                serde_json::json!({"tokens": count}),
            ));
        }

        let path = args.get_arg(0).unwrap();
        let path_buf = PathBuf::from(path);

        if !path_buf.exists() {
            return Ok(ToolResult::error(format!("Path not found: {}", path)));
        }

        let content = fs::read_to_string(&path_buf)
            .map_err(|e| anyhow::anyhow!("Failed to read file: {}", e))?;

        let count = Self::count_from_text(&content);

        Ok(ToolResult::success_with_data(
            format!("Total tokens: {}", count),
            serde_json::json!({"tokens": count, "path": path}),
        ))
    }

    fn get_parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "file_path": { "type": "string", "description": "Path to file to count tokens for (positional)" },
                "content": { "type": "string", "description": "Direct content to count tokens for (named)" }
            },
            "required": []
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::ToolArgs;
    use crate::tools::old::state::ToolState;
    use std::sync::{Arc, Mutex};

    #[test]
    fn test_count_tokens_from_text() {
        let mut tool = CountTokensTool::new();
        let args = ToolArgs::from_args(&["--content=hello world"]);
        let state = Arc::new(Mutex::new(ToolState::new()));
        let res = tool.execute(&args, &state).unwrap();
        assert!(res.message.contains("Total tokens:"));
    }

    #[test]
    fn test_count_tokens_file() {
        let tmp = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(tmp.path(), "This is a small test file").unwrap();
        let mut tool = CountTokensTool::new();
        let args = ToolArgs::from_args(&[tmp.path().to_string_lossy().as_ref()]);
        let state = Arc::new(Mutex::new(ToolState::new()));
        let res = tool.execute(&args, &state).unwrap();
        assert!(res.message.contains("Total tokens:"));
    }
}