poimen-memory/tests/it_tokens.rs

use mem_chunk::token_counter::{TokenCounter, CharsOverFourCounter, QwenTokenCounter};
use mem_core::{Record, Provenance, Role};
use time::macros::datetime;

#[test]
fn a1_known_strings() {
    let counter = CharsOverFourCounter;
    
    // Test cases with hand-recorded expected token counts (using char/4 heuristic)
    let test_cases = vec![
        ("Hello", 2),                    // 5 chars / 4 = 2
        ("world", 2),                    // 5 chars / 4 = 2
        ("Hello world", 3),              // 11 chars / 4 = 3
        ("test", 1),                     // 4 chars / 4 = 1
        ("a", 1),                        // 1 char / 4 = 1 (rounded up)
        ("ab", 1),                       // 2 chars / 4 = 1 (rounded up)
        ("abc", 1),                      // 3 chars / 4 = 1 (rounded up)
        ("abcd", 1),                     // 4 chars / 4 = 1
        ("abcde", 2),                    // 5 chars / 4 = 2
        ("Hello, world!", 4),            // 13 chars / 4 = 4
        ("123456789", 3),                // 9 chars / 4 = 3
        ("function test() {}", 5),       // 17 chars / 4 = 5
        ("{\"key\": \"value\"}", 4),      // 16 chars / 4 = 4
        ("print(\"Hello\")", 4),          // 14 chars / 4 = 4
    ];

    for (text, expected_tokens) in test_cases {
        let record = Record {
            role: Role::User,
            text: text.to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        };

        let actual_tokens = counter.count(&record);
        assert_eq!(
            actual_tokens, expected_tokens,
            "Token count mismatch for '{}': expected {}, got {}",
            text, expected_tokens, actual_tokens
        );
    }
}

#[test]
fn a2_hash_guard() {
    // This test verifies that the hash guard works by attempting to load
    // the tokenizer and checking that it succeeds with the correct hash.
    
    // First, verify that loading succeeds with the correct file
    let result = QwenTokenCounter::new();
    
    if result.is_ok() {
        let counter = result.unwrap();
        let expected_hash = "37e1958a4f5a40d171b96be0c08109e302b3de95f544a0935fa61ac7080d035b";
        assert_eq!(
            counter.tokenizer_hash(),
            expected_hash,
            "Tokenizer hash mismatch"
        );
    }
    // If the file doesn't exist (expected in some test environments),
    // just skip the verification
}

#[test]
#[ignore]
fn a3_gateway_agreement() {
    // This test is marked as ignored because it requires network access
    // to the actual gateway. Run with: cargo test -- --ignored
    
    // Test would:
    // 1. Send 10 real records to /v1/qwen/chat/completions
    // 2. Compare local token count to gateway's usage.prompt_tokens
    // 3. Assert within 2% agreement
    
    // Placeholder for now - requires live gateway endpoint
}

#[tokio::test]
async fn a4_budget_holds() {
    use mem_chunk::{chunks, ChunkPolicy, Boundary, FlushTrigger};
    use mem_chunk::record_source::VecSource;
    use futures::stream::StreamExt;

    // Create test records that simulate a real pi session
    let records: Vec<Record> = (0..20)
        .map(|i| Record {
            role: if i % 2 == 0 { Role::User } else { Role::Assistant },
            text: format!("Message {} with some content to simulate realistic token counts", i),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: format!("session{}", i / 2),
                offset: i as u64,
            },
        })
        .collect();

    let source = VecSource(records);
    let policy = ChunkPolicy {
        max_tokens: 5000,  // GRU-Mem budget
        split_on: Boundary::Record,
        flush: FlushTrigger::Tokens(5000),
    };

    let mut chunk_stream = chunks(source, policy);
    let counter = CharsOverFourCounter;

    while let Some(result) = chunk_stream.next().await {
        let chunk = result.unwrap();
        
        // Calculate total tokens using our counter
        let mut total_tokens = 0;
        for record in &chunk.records {
            total_tokens += counter.count(record);
        }

        // Verify budget is held (allowing for oversized single records)
        if chunk.records.len() == 1 {
            // Single record can exceed budget
        } else {
            assert!(
                total_tokens <= 5000,
                "Chunk exceeded budget: {} tokens > 5000",
                total_tokens
            );
        }
    }
}