#![cfg(feature = "net")]
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use llmp::tokenize::enrich;
use llmp::{Config, Record, Source, analyze};
use serde_json::{Value, json};
#[derive(Clone, Copy)]
enum Serves {
LlamaCpp,
Ollama,
Nothing,
}
struct Fake {
addr: SocketAddr,
seen: Arc<Mutex<Vec<(String, Value)>>>,
}
impl Fake {
fn url(&self) -> String {
format!("http://{}", self.addr)
}
fn seen(&self) -> Vec<(String, Value)> {
self.seen.lock().unwrap().clone()
}
}
fn tokens_for(text: &str) -> usize {
text.chars().count().div_ceil(3)
}
fn fake(serves: Serves) -> Fake {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let seen = Arc::new(Mutex::new(Vec::new()));
let log = Arc::clone(&seen);
std::thread::spawn(move || {
for stream in listener.incoming() {
let Ok(stream) = stream else { return };
let log = Arc::clone(&log);
std::thread::spawn(move || serve(stream, serves, log));
}
});
Fake { addr, seen }
}
fn serve(stream: TcpStream, serves: Serves, log: Arc<Mutex<Vec<(String, Value)>>>) {
let mut reader = BufReader::new(stream.try_clone().unwrap());
let mut writer = stream;
loop {
let mut first = String::new();
if reader.read_line(&mut first).unwrap_or(0) == 0 {
return;
}
let path = first.split_whitespace().nth(1).unwrap_or("").to_string();
let mut length = 0;
loop {
let mut line = String::new();
reader.read_line(&mut line).unwrap();
if line == "\r\n" {
break;
}
if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
length = v.trim().parse().unwrap();
}
}
let mut body = vec![0; length];
reader.read_exact(&mut body).unwrap();
let body: Value = serde_json::from_slice(&body).unwrap_or(Value::Null);
log.lock().unwrap().push((path.clone(), body.clone()));
let reply = match (serves, path.as_str()) {
(Serves::LlamaCpp, "/tokenize") => Some(body["content"].as_str().map(tokens_for)),
(Serves::Ollama, "/api/tokenize") => Some(body["text"].as_str().map(tokens_for)),
_ => None,
};
let (status, text) = match reply {
Some(Some(n)) => ("200 OK", json!({"tokens": vec![7; n]}).to_string()),
Some(None) => ("400 Bad Request", "{}".to_string()),
None => ("404 Not Found", "{\"error\":\"not found\"}".to_string()),
};
let response = format!(
"HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{text}",
text.len()
);
if writer.write_all(response.as_bytes()).is_err() {
return;
}
}
}
const SYSTEM: &str = "You are a careful coding agent. Keep answers short.";
fn tool() -> Value {
json!({"type": "function", "function": {"name": "read_file", "parameters": {"type": "object"}}})
}
fn chat(id: u64, messages: Value) -> Record {
serde_json::from_value(json!({
"v": 1, "id": id, "conn": 1, "t_start_ms": id * 1000, "t_first_byte_ms": null, "t_end_ms": id * 1000 + 500,
"request": {"method": "POST", "path": "/v1/chat/completions", "headers": {}, "body_bytes": 1,
"body": {"model": "m", "messages": messages, "tools": [tool()]}, "body_text": null, "digest": null},
"response": {"status": 200, "headers": {}, "stream": null, "chunks": null, "bad_chunks": null, "body_bytes": 1,
"truncated": false, "body": {"choices": [{"message": {"role": "assistant", "content": "ok"}}]},
"body_text": null, "digest": null},
"error": null, "tokens": null
}))
.unwrap()
}
fn records() -> Vec<Record> {
vec![
chat(
1,
json!([{"role": "system", "content": SYSTEM}, {"role": "user", "content": "Read main.rs please."}]),
),
chat(
2,
json!([
{"role": "system", "content": SYSTEM},
{"role": "user", "content": "Read main.rs please."},
{"role": "assistant", "content": null},
{"role": "user", "content": "Now lib.rs."}
]),
),
]
}
fn exact_rows(records: &[Record]) -> bool {
analyze(records, &Config::default())
.rows
.iter()
.all(|r| r.composition.system.source == Source::Tokenizer)
}
#[test]
fn llama_cpp_counts_every_distinct_text_once() {
let server = fake(Serves::LlamaCpp);
let mut records = records();
let result = enrich(&mut records, &server.url());
assert_eq!(result.warnings, Vec::<String>::new());
assert_eq!(result.counted, 5);
let seen = server.seen();
assert_eq!(seen.len(), 4, "empty texts are not sent: {seen:?}");
assert!(
seen.iter()
.all(|(path, body)| path == "/tokenize" && body["add_special"] == false)
);
let tokens = records[1].tokens.as_ref().unwrap();
assert_eq!(tokens.messages.len(), 4);
assert_eq!(tokens.tools.len(), 1);
assert!(tokens.messages.values().any(|&n| n == 0));
assert!(records[0].tokens.as_ref().unwrap().messages.len() == 2);
let report = analyze(&records, &Config::default());
let first = &report.rows[0].composition;
assert_eq!(first.system.source, Source::Tokenizer);
assert_eq!(first.system.value as usize, tokens_for(SYSTEM));
assert_eq!(
first.new_content.value as usize,
tokens_for("Read main.rs please.")
);
let tool_text = serde_json::to_string(&tool()).unwrap();
assert_eq!(first.tools.value as usize, tokens_for(&tool_text));
}
#[test]
fn a_second_run_sends_nothing() {
let server = fake(Serves::LlamaCpp);
let mut records = records();
enrich(&mut records, &server.url());
let before = records.clone();
let sent = server.seen().len();
let again = enrich(&mut records, &server.url());
assert_eq!(again.counted, 0);
assert_eq!(server.seen().len(), sent);
assert_eq!(records, before);
}
#[test]
fn ollama_endpoint_is_the_fallback() {
let server = fake(Serves::Ollama);
let mut records = records();
let result = enrich(&mut records, &server.url());
assert_eq!(result.warnings, Vec::<String>::new());
assert_eq!(result.counted, 5);
let seen = server.seen();
assert_eq!(seen[0].0, "/tokenize");
assert!(
seen[1..]
.iter()
.all(|(path, body)| path == "/api/tokenize" && body["model"] == "m")
);
assert!(exact_rows(&records));
}
#[test]
fn no_endpoint_keeps_the_estimates() {
let server = fake(Serves::Nothing);
let mut records = records();
let before = records.clone();
let result = enrich(&mut records, &server.url());
assert_eq!(result.counted, 0);
assert_eq!(result.warnings.len(), 1);
assert!(
result.warnings[0].contains("no tokenize endpoint"),
"{:?}",
result.warnings
);
assert!(result.warnings[0].contains("/tokenize answered 404"));
assert_eq!(records, before);
let report = analyze(&records, &Config::default());
assert!(
report
.rows
.iter()
.all(|r| r.composition.system.source == Source::Estimate)
);
}
#[test]
fn unreachable_server_keeps_the_estimates() {
let port = TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port();
let mut records = records();
let before = records.clone();
let result = enrich(&mut records, &format!("http://127.0.0.1:{port}"));
assert_eq!(result.counted, 0);
assert!(
result.warnings[0].contains("cannot reach"),
"{:?}",
result.warnings
);
assert_eq!(records, before);
}
#[test]
fn bad_urls_are_warnings() {
let mut records = records();
for url in ["https://example.com", "127.0.0.1:8080", "::"] {
let result = enrich(&mut records, url);
assert_eq!(result.counted, 0);
assert_eq!(result.warnings.len(), 1, "{url}: {:?}", result.warnings);
}
assert!(records.iter().all(|r| r.tokens.is_none()));
}
#[test]
fn digest_records_are_skipped_with_a_warning() {
let server = fake(Serves::LlamaCpp);
let mut records = records();
records[0].request.body = None;
records[0].request.digest = Some(serde_json::from_value(json!({
"model": "m", "params": {"max_tokens": null, "response_format": false, "grammar": false, "stream": false},
"messages": [], "tools": [], "tools_hash": null
})).unwrap());
let result = enrich(&mut records, &server.url());
assert!(
result
.warnings
.iter()
.any(|w| w.contains("1 records were recorded in digest mode"))
);
assert_eq!(result.counted, 5);
assert!(records[0].tokens.is_none());
assert!(records[1].tokens.is_some());
}
#[cfg(feature = "cli")]
#[test]
fn cli_tokenize_flag() {
let server = fake(Serves::LlamaCpp);
let path = std::path::PathBuf::from(env!("CARGO_TARGET_TMPDIR")).join("tokenize_cli.jsonl");
let text: String = records()
.iter()
.map(|r| format!("{}\n", serde_json::to_string(r).unwrap()))
.collect();
std::fs::write(&path, text).unwrap();
let run = |url: &str| {
std::process::Command::new(env!("CARGO_BIN_EXE_llmp"))
.args([
"analyze",
path.to_str().unwrap(),
"--json",
"--tokenize",
url,
])
.output()
.unwrap()
};
let out = run(&server.url());
assert_eq!(out.status.code(), Some(0));
assert!(String::from_utf8_lossy(&out.stderr).contains("counted tokens for 5 texts"));
let report: llmp::Report = serde_json::from_slice(&out.stdout).unwrap();
assert!(
report
.rows
.iter()
.all(|r| r.composition.system.source == Source::Tokenizer)
);
let port = TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port();
let out = run(&format!("http://127.0.0.1:{port}"));
assert_eq!(out.status.code(), Some(0));
assert!(String::from_utf8_lossy(&out.stderr).contains("warning: tokenize: cannot reach"));
}