poimen-memory/crates/mem-chunk/src/token_counter.rs

use mem_core::Record;
use sha2::{Digest, Sha256};

/// Token counter trait.
pub trait TokenCounter {
    /// Count tokens in a record.
    fn count(&self, record: &Record) -> usize;
}

/// Stub token counter: characters / 4
/// Simple heuristic for testing; real counter uses a proper tokenizer.
#[derive(Debug, Clone)]
pub struct CharsOverFourCounter;

impl TokenCounter for CharsOverFourCounter {
    fn count(&self, record: &Record) -> usize {
        // Rough heuristic: 4 characters per token
        record.text.len().div_ceil(4)
    }
}

/// Qwen2 BPE tokenizer-backed token counter.
/// Uses the vendored tokenizer.json with hash verification.
pub struct QwenTokenCounter {
    tokenizer: tokenizers::Tokenizer,
    tokenizer_hash: String,
}

impl QwenTokenCounter {
    /// Load the Qwen2 tokenizer from the vendored file.
    /// Returns an error if the file hash doesn't match the expected value.
    pub fn new() -> anyhow::Result<Self> {
        const EXPECTED_HASH: &str = "37e1958a4f5a40d171b96be0c08109e302b3de95f544a0935fa61ac7080d035b";
        const TOKENIZER_PATH: &str = "assets/qwen2-tokenizer.json";

        // Read and verify the tokenizer file hash
        let tokenizer_bytes = std::fs::read(TOKENIZER_PATH)
            .map_err(|e| anyhow::anyhow!("Failed to read {}: {}", TOKENIZER_PATH, e))?;

        let mut hasher = Sha256::new();
        hasher.update(&tokenizer_bytes);
        let hash = hasher.finalize();
        let hash_hex = hex::encode(hash);

        if hash_hex != EXPECTED_HASH {
            return Err(anyhow::anyhow!(
                "Tokenizer hash mismatch for {}: expected {}, got {}",
                TOKENIZER_PATH,
                EXPECTED_HASH,
                hash_hex
            ));
        }

        let tokenizer = tokenizers::Tokenizer::from_bytes(&tokenizer_bytes)
            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;

        Ok(QwenTokenCounter {
            tokenizer,
            tokenizer_hash: hash_hex,
        })
    }

    /// Get the hash of the loaded tokenizer
    pub fn tokenizer_hash(&self) -> &str {
        &self.tokenizer_hash
    }
}

impl TokenCounter for QwenTokenCounter {
    fn count(&self, record: &Record) -> usize {
        // Tokenize the text and count tokens
        match self.tokenizer.encode(record.text.as_str(), false) {
            Ok(encoding) => encoding.get_tokens().len(),
            Err(_) => {
                // Fallback to character-based estimate if tokenization fails
                record.text.len().div_ceil(4)
            }
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use mem_core::{Provenance, Role};
    use time::macros::datetime;

    #[test]
    fn test_chars_over_four_counter() {
        let counter = CharsOverFourCounter;

        let record = Record {
            role: Role::User,
            text: "Hello".to_string(),  // 5 chars = 2 tokens (rounded up)
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        };

        assert_eq!(counter.count(&record), 2);
    }

    #[test]
    fn test_qwen_token_counter_loads() {
        let result = QwenTokenCounter::new();
        // This test will pass if the tokenizer loads successfully
        // or fail if the file doesn't exist or hash mismatches
        if result.is_ok() {
            let counter = result.unwrap();
            assert!(!counter.tokenizer_hash.is_empty());
        }
    }
}