132 lines
4.5 KiB
Plaintext
132 lines
4.5 KiB
Plaintext
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
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|