use std::collections::{BTreeMap, BTreeSet};
use std::time::Duration;
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::client::conn::http1::{self, SendRequest};
use hyper::{Method, Request, Uri, header};
use hyper_util::rt::TokioIo;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tokio::net::TcpStream;
use crate::hash;
use crate::normalize::{self, Api};
use crate::record::{self, Record, TokenCounts};
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct TokenizeResult {
pub counted: u64,
pub warnings: Vec<String>,
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum Kind {
Message,
Tool,
}
struct Text {
kind: Kind,
hash: String,
text: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Flavor {
LlamaCpp,
Ollama,
}
#[derive(Debug)]
struct Endpoint {
socket: String,
authority: String,
base: String,
}
pub fn enrich(records: &mut [Record], server: &str) -> TokenizeResult {
let mut result = TokenizeResult::default();
let endpoint = match Endpoint::parse(server) {
Ok(endpoint) => endpoint,
Err(warning) => {
result.warnings.push(warning);
return result;
}
};
let (texts, models, digest_records) = collect(records);
if digest_records > 0 {
result.warnings.push(format!(
"{digest_records} records were recorded in digest mode and have no text to tokenize"
));
}
if models.len() > 1 {
result.warnings.push(format!(
"the recording uses {} models; every text was counted with the tokenizer of the model {server} serves",
models.len()
));
}
if texts.is_empty() {
return result;
}
let model = models.into_iter().next().unwrap_or_default();
let run =
std::thread::scope(|scope| scope.spawn(|| count_all(&endpoint, &model, &texts)).join());
let (counts, warnings) = run.unwrap_or_else(|_| {
(
BTreeMap::new(),
vec!["the tokenize thread panicked".to_string()],
)
});
result.warnings.extend(warnings);
result.counted = counts.len() as u64;
apply(records, &counts);
result
}
impl Endpoint {
fn parse(server: &str) -> Result<Endpoint, String> {
let uri: Uri = server
.parse()
.map_err(|e| format!("--tokenize {server}: not a URL: {e}"))?;
match uri.scheme_str() {
Some("http") => {}
Some(other) => {
return Err(format!(
"--tokenize {server}: only http:// is supported, not {other}://"
));
}
None => return Err(format!("--tokenize {server}: the URL needs http://")),
}
let authority = uri
.authority()
.ok_or_else(|| format!("--tokenize {server}: the URL has no host"))?;
let port = authority.port_u16().unwrap_or(80);
Ok(Endpoint {
socket: format!("{}:{port}", authority.host()),
authority: authority.to_string(),
base: uri.path().trim_end_matches('/').to_string(),
})
}
}
fn collect(records: &[Record]) -> (Vec<Text>, BTreeSet<String>, usize) {
let mut texts = Vec::new();
let mut seen = BTreeSet::new();
let mut models = BTreeSet::new();
let mut digest_records = 0;
for record in records {
let api = normalize::detect_api(&record.request.path);
if api == Api::Other {
continue;
}
if record.request.digest.is_some() {
digest_records += 1;
continue;
}
let Some(body) = &record.request.body else {
continue;
};
let Some(digest) = record::request_digest(api, body) else {
continue;
};
if let Some(model) = digest.model {
models.insert(model);
}
let known = record.tokens.as_ref();
let messages = body.get("messages").and_then(Value::as_array);
for (message, value) in digest.messages.iter().zip(messages.into_iter().flatten()) {
let counted = known.is_some_and(|t| t.messages.contains_key(&message.hash));
if !counted && seen.insert((Kind::Message, message.hash.clone())) {
texts.push(Text {
kind: Kind::Message,
hash: message.hash.clone(),
text: record::message_text(value),
});
}
}
let tools = body.get("tools").and_then(Value::as_array);
for (tool, value) in digest.tools.iter().zip(tools.into_iter().flatten()) {
let counted = known.is_some_and(|t| t.tools.contains_key(&tool.hash));
if !counted && seen.insert((Kind::Tool, tool.hash.clone())) {
texts.push(Text {
kind: Kind::Tool,
hash: tool.hash.clone(),
text: hash::canonical_json(value),
});
}
}
}
(texts, models, digest_records)
}
fn apply(records: &mut [Record], counts: &Counts) {
if counts.is_empty() {
return;
}
for record in records.iter_mut() {
let api = normalize::detect_api(&record.request.path);
let Some(digest) = record
.request
.body
.as_ref()
.filter(|_| api != Api::Other)
.and_then(|body| record::request_digest(api, body))
else {
continue;
};
let mut added = TokenCounts::default();
for message in &digest.messages {
if let Some(&n) = counts.get(&(Kind::Message, message.hash.clone())) {
added.messages.insert(message.hash.clone(), n);
}
}
for tool in &digest.tools {
if let Some(&n) = counts.get(&(Kind::Tool, tool.hash.clone())) {
added.tools.insert(tool.hash.clone(), n);
}
}
if added.messages.is_empty() && added.tools.is_empty() {
continue;
}
let tokens = record.tokens.get_or_insert_with(TokenCounts::default);
tokens.messages.extend(added.messages);
tokens.tools.extend(added.tools);
}
}
type Counts = BTreeMap<(Kind, String), u32>;
fn count_all(endpoint: &Endpoint, model: &str, texts: &[Text]) -> (Counts, Vec<String>) {
let runtime = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(runtime) => runtime,
Err(e) => {
return (
Counts::new(),
vec![format!("cannot start the tokenize runtime: {e}")],
);
}
};
runtime.block_on(async {
let mut counts = Counts::new();
let mut warnings = Vec::new();
let mut client = Client {
endpoint,
sender: None,
};
let mut flavor = None;
let mut failed = 0u64;
for text in texts {
let key = (text.kind, text.hash.clone());
if text.text.is_empty() {
counts.insert(key, 0);
continue;
}
let flavor = match flavor {
Some(flavor) => flavor,
None => match detect(&mut client, model, &text.text).await {
Ok((detected, n)) => {
flavor = Some(detected);
counts.insert(key, n);
continue;
}
Err(warning) => {
warnings.push(warning);
return (counts, warnings);
}
},
};
match count(&mut client, flavor, model, &text.text).await {
Ok(Some(n)) => {
counts.insert(key, n);
}
Ok(None) => failed += 1,
Err(warning) => {
warnings.push(format!(
"{warning}; {} texts were not counted",
texts.len() - counts.len()
));
return (counts, warnings);
}
}
}
if failed > 0 {
warnings.push(format!(
"{failed} texts could not be tokenized and keep their estimate"
));
}
(counts, warnings)
})
}
async fn detect(client: &mut Client<'_>, model: &str, text: &str) -> Result<(Flavor, u32), String> {
let socket = client.endpoint.socket.clone();
let mut statuses = Vec::new();
for flavor in [Flavor::LlamaCpp, Flavor::Ollama] {
let (path, body) = request(client.endpoint, flavor, model, text);
let (status, reply) = client.post(&path, &body).await.map_err(|e| {
format!("cannot reach the tokenize endpoint at {socket}: {e}; estimates are kept")
})?;
if let Some(n) = tokens(status, &reply) {
return Ok((flavor, n));
}
statuses.push(format!("{path} answered {status}"));
}
Err(format!(
"no tokenize endpoint at {socket} ({}); estimates are kept",
statuses.join(", ")
))
}
async fn count(
client: &mut Client<'_>,
flavor: Flavor,
model: &str,
text: &str,
) -> Result<Option<u32>, String> {
let (path, body) = request(client.endpoint, flavor, model, text);
let (status, reply) = client.post(&path, &body).await.map_err(|e| {
format!(
"lost the tokenize endpoint at {}: {e}",
client.endpoint.socket
)
})?;
Ok(tokens(status, &reply))
}
fn request(endpoint: &Endpoint, flavor: Flavor, model: &str, text: &str) -> (String, Value) {
match flavor {
Flavor::LlamaCpp => (
format!("{}/tokenize", endpoint.base),
json!({ "content": text, "add_special": false }),
),
Flavor::Ollama => (
format!("{}/api/tokenize", endpoint.base),
json!({ "model": model, "text": text }),
),
}
}
fn tokens(status: u16, reply: &Value) -> Option<u32> {
if !(200..300).contains(&status) {
return None;
}
let tokens = reply.get("tokens")?.as_array()?;
u32::try_from(tokens.len()).ok()
}
struct Client<'a> {
endpoint: &'a Endpoint,
sender: Option<SendRequest<Full<Bytes>>>,
}
impl Client<'_> {
async fn post(&mut self, path: &str, body: &Value) -> Result<(u16, Value), String> {
let bytes = Bytes::from(body.to_string());
for attempt in 0..2 {
let reused = self.sender.as_ref().is_some_and(|s| !s.is_closed());
if !reused {
self.sender = Some(self.connect().await?);
}
let Some(sender) = self.sender.as_mut() else {
return Err("no connection".into());
};
let request = Request::builder()
.method(Method::POST)
.uri(path)
.header(header::HOST, &self.endpoint.authority)
.header(header::CONTENT_TYPE, "application/json")
.body(Full::new(bytes.clone()))
.map_err(|e| e.to_string())?;
let sent = async {
sender.ready().await?;
let response = sender.send_request(request).await?;
let status = response.status().as_u16();
let body = response.into_body().collect().await?.to_bytes();
Ok::<_, hyper::Error>((status, body))
};
match tokio::time::timeout(REQUEST_TIMEOUT, sent).await {
Ok(Ok((status, body))) => {
let reply = serde_json::from_slice(&body).unwrap_or(Value::Null);
return Ok((status, reply));
}
Ok(Err(e)) => {
self.sender = None;
if reused && attempt == 0 {
continue;
}
return Err(e.to_string());
}
Err(_) => {
self.sender = None;
return Err(format!("no answer within {} s", REQUEST_TIMEOUT.as_secs()));
}
}
}
Err("connection closed".into())
}
async fn connect(&self) -> Result<SendRequest<Full<Bytes>>, String> {
let stream =
tokio::time::timeout(CONNECT_TIMEOUT, TcpStream::connect(&self.endpoint.socket))
.await
.map_err(|_| format!("connecting took more than {} s", CONNECT_TIMEOUT.as_secs()))?
.map_err(|e| e.to_string())?;
let _ = stream.set_nodelay(true);
let (sender, conn) = http1::handshake(TokioIo::new(stream))
.await
.map_err(|e| e.to_string())?;
tokio::spawn(async move {
let _ = conn.await;
});
Ok(sender)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn endpoint_urls() {
let e = Endpoint::parse("http://127.0.0.1:8080").unwrap();
assert_eq!(
(e.socket.as_str(), e.authority.as_str(), e.base.as_str()),
("127.0.0.1:8080", "127.0.0.1:8080", "")
);
let e = Endpoint::parse("http://localhost/llama/").unwrap();
assert_eq!(
(e.socket.as_str(), e.base.as_str()),
("localhost:80", "/llama")
);
assert!(
Endpoint::parse("https://example.com")
.unwrap_err()
.contains("only http://")
);
assert!(Endpoint::parse("127.0.0.1:8080").is_err());
assert!(Endpoint::parse("not a url").is_err());
}
#[test]
fn token_replies() {
assert_eq!(tokens(200, &json!({"tokens": [1, 2, 3]})), Some(3));
assert_eq!(tokens(200, &json!({"tokens": []})), Some(0));
assert_eq!(tokens(404, &json!({"tokens": [1]})), None);
assert_eq!(tokens(200, &json!({"error": "x"})), None);
assert_eq!(tokens(200, &Value::Null), None);
}
}