rskim 2.3.1

The most intelligent context optimization engine for coding agents. Code-aware AST parsing, command rewriting, output compression.
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

//! Token counting using OpenAI's tiktoken tokenizer
//!
//! ARCHITECTURE: Uses cl100k_base encoding (GPT-3.5-turbo, GPT-4)
//! - Provides accurate token counts for LLM context window calculation
//! - Counts tokens before and after transformation
//! - Reports reduction statistics
//! - Lazy initialization to avoid recreating tokenizer on every call

use anyhow::Result;
use std::sync::OnceLock;
use tiktoken_rs::{cl100k_base, CoreBPE};

/// Global tokenizer instance (lazy-initialized on first use)
static TOKENIZER: OnceLock<CoreBPE> = OnceLock::new();

/// Get or initialize the global tokenizer instance
fn get_tokenizer() -> &'static CoreBPE {
    TOKENIZER.get_or_init(|| cl100k_base().expect("Failed to initialize cl100k_base tokenizer"))
}

/// Count tokens in text using cl100k_base encoding (GPT-3.5-turbo, GPT-4)
pub(crate) fn count_tokens(text: &str) -> Result<usize> {
    let tokenizer = get_tokenizer();
    let tokens = tokenizer.encode_with_special_tokens(text);
    Ok(tokens.len())
}

/// Statistics for token reduction
#[derive(Debug, Clone)]
pub(crate) struct TokenStats {
    /// Original token count
    pub(crate) original: usize,
    /// Transformed token count
    pub(crate) transformed: usize,
}

impl TokenStats {
    /// Create new token stats
    pub(crate) fn new(original: usize, transformed: usize) -> Self {
        Self {
            original,
            transformed,
        }
    }

    /// Calculate reduction percentage (negative if transformed is larger)
    pub(crate) fn reduction_percentage(&self) -> f32 {
        if self.original == 0 {
            return 0.0;
        }
        ((self.original as f32 - self.transformed as f32) / self.original as f32) * 100.0
    }

    /// Get tokens saved
    #[allow(dead_code)]
    pub(crate) fn tokens_saved(&self) -> usize {
        self.original.saturating_sub(self.transformed)
    }

    /// Format stats for display
    pub(crate) fn format(&self) -> String {
        format!(
            "{} tokens → {} tokens ({:.1}% reduction)",
            format_number(self.original),
            format_number(self.transformed),
            self.reduction_percentage()
        )
    }
}

/// Format number with thousands separators
pub(crate) fn format_number(n: usize) -> String {
    let s = n.to_string();
    let mut result = String::new();

    for (count, ch) in s.chars().rev().enumerate() {
        if count > 0 && count % 3 == 0 {
            result.push(',');
        }
        result.push(ch);
    }

    result.chars().rev().collect()
}

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

    #[test]
    fn test_count_tokens() {
        let text = "Hello, world!";
        let count = count_tokens(text).unwrap();
        assert!(count > 0);
        assert!(count < 10); // Should be around 3-4 tokens
    }

    #[test]
    fn test_token_stats() {
        let stats = TokenStats::new(1000, 200);
        assert_eq!(stats.reduction_percentage(), 80.0);
        assert_eq!(stats.tokens_saved(), 800);
    }

    #[test]
    fn test_format_number() {
        assert_eq!(format_number(1000), "1,000");
        assert_eq!(format_number(1000000), "1,000,000");
        assert_eq!(format_number(123), "123");
    }

    #[test]
    fn test_stats_format() {
        let stats = TokenStats::new(1000, 200);
        let formatted = stats.format();
        assert!(formatted.contains("1,000"));
        assert!(formatted.contains("200"));
        assert!(formatted.contains("80.0%"));
    }
}