oxideshield_core/

tokenizer.rs

//! Tokenizer integration using tiktoken-rs

use crate::{Error, Result};
use tiktoken_rs::{cl100k_base, o200k_base, p50k_base, CoreBPE};
use tracing::debug;

/// Supported tokenizer models
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum TokenizerModel {
    /// GPT-4, GPT-3.5-turbo, text-embedding-ada-002
    #[default]
    Cl100kBase,
    /// GPT-4o, GPT-4o-mini
    O200kBase,
    /// GPT-3 models (text-davinci-003, etc.)
    P50kBase,
}

impl std::str::FromStr for TokenizerModel {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "cl100k_base" | "cl100k" | "gpt-4" | "gpt-3.5-turbo" => Ok(TokenizerModel::Cl100kBase),
            "o200k_base" | "o200k" | "gpt-4o" => Ok(TokenizerModel::O200kBase),
            "p50k_base" | "p50k" | "gpt-3" => Ok(TokenizerModel::P50kBase),
            _ => Err(format!("Unknown tokenizer model: {}", s)),
        }
    }
}

/// Tokenizer for counting and encoding tokens
pub struct Tokenizer {
    bpe: CoreBPE,
    model: TokenizerModel,
}

impl std::fmt::Debug for Tokenizer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Tokenizer")
            .field("model", &self.model)
            .finish()
    }
}

impl Tokenizer {
    /// Create a new tokenizer with the specified model
    pub fn new(model: TokenizerModel) -> Result<Self> {
        let bpe = match model {
            TokenizerModel::Cl100kBase => {
                cl100k_base().map_err(|e| Error::Tokenizer(e.to_string()))?
            }
            TokenizerModel::O200kBase => {
                o200k_base().map_err(|e| Error::Tokenizer(e.to_string()))?
            }
            TokenizerModel::P50kBase => p50k_base().map_err(|e| Error::Tokenizer(e.to_string()))?,
        };

        debug!("Created tokenizer with model {:?}", model);

        Ok(Self { bpe, model })
    }

    /// Create a tokenizer with the default model (cl100k_base)
    pub fn default_model() -> Result<Self> {
        Self::new(TokenizerModel::default())
    }

    /// Get the tokenizer model
    pub fn model(&self) -> TokenizerModel {
        self.model
    }

    /// Count the number of tokens in the text
    pub fn count_tokens(&self, text: &str) -> usize {
        self.bpe.encode_ordinary(text).len()
    }

    /// Encode text into token IDs
    pub fn encode(&self, text: &str) -> Vec<u32> {
        self.bpe.encode_ordinary(text)
    }

    /// Decode token IDs back to text
    pub fn decode(&self, tokens: &[u32]) -> Result<String> {
        self.bpe
            .decode(tokens.to_vec())
            .map_err(|e| Error::Tokenizer(e.to_string()))
    }

    /// Truncate text to a maximum number of tokens
    pub fn truncate(&self, text: &str, max_tokens: usize) -> Result<String> {
        let tokens = self.encode(text);
        if tokens.len() <= max_tokens {
            return Ok(text.to_string());
        }

        let truncated = &tokens[..max_tokens];
        self.decode(truncated)
    }

    /// Check if text exceeds a token limit
    pub fn exceeds_limit(&self, text: &str, limit: usize) -> bool {
        self.count_tokens(text) > limit
    }
}

impl Default for Tokenizer {
    fn default() -> Self {
        Self::new(TokenizerModel::default()).expect("Failed to create default tokenizer")
    }
}

/// Token statistics for text analysis
#[derive(Debug, Clone, Default)]
pub struct TokenStats {
    /// Total token count
    pub total_tokens: usize,
    /// Average tokens per line
    pub avg_tokens_per_line: f64,
    /// Maximum tokens in a single line
    pub max_line_tokens: usize,
    /// Number of lines
    pub line_count: usize,
}

impl Tokenizer {
    /// Calculate token statistics for text
    pub fn stats(&self, text: &str) -> TokenStats {
        let lines: Vec<&str> = text.lines().collect();
        let line_count = lines.len();

        if line_count == 0 {
            return TokenStats::default();
        }

        let mut total_tokens = 0;
        let mut max_line_tokens = 0;

        for line in &lines {
            let tokens = self.count_tokens(line);
            total_tokens += tokens;
            max_line_tokens = max_line_tokens.max(tokens);
        }

        TokenStats {
            total_tokens,
            avg_tokens_per_line: total_tokens as f64 / line_count as f64,
            max_line_tokens,
            line_count,
        }
    }
}

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

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

    #[test]
    fn test_encode_decode() {
        let tokenizer = Tokenizer::default();
        let text = "Hello, world!";
        let tokens = tokenizer.encode(text);
        let decoded = tokenizer.decode(&tokens).unwrap();
        assert_eq!(decoded, text);
    }

    #[test]
    fn test_truncate() {
        let tokenizer = Tokenizer::default();
        let text = "This is a longer text that should be truncated to fewer tokens.";
        let truncated = tokenizer.truncate(text, 5).unwrap();
        assert!(tokenizer.count_tokens(&truncated) <= 5);
    }

    #[test]
    fn test_exceeds_limit() {
        let tokenizer = Tokenizer::default();
        assert!(!tokenizer.exceeds_limit("short", 100));
        assert!(tokenizer.exceeds_limit("This is a longer text.", 2));
    }

    #[test]
    fn test_stats() {
        let tokenizer = Tokenizer::default();
        let text = "Line one\nLine two\nLine three";
        let stats = tokenizer.stats(text);

        assert_eq!(stats.line_count, 3);
        assert!(stats.total_tokens > 0);
    }
}