use std::sync::OnceLock;
use polyc_llm::Usage;
use prometheus::{IntCounterVec, register_int_counter_vec};
fn turn_prompt_tokens() -> &'static IntCounterVec {
static V: OnceLock<IntCounterVec> = OnceLock::new();
V.get_or_init(|| {
register_int_counter_vec!(
"polychrome_turn_prompt_tokens_total",
"Per-turn prompt token counts by kind (input, cache_read, cache_creation); \
prompt-cache effectiveness is cache_read / input.",
&["token_kind"]
)
.expect("register polychrome_turn_prompt_tokens_total")
})
}
pub(crate) fn record_turn(usage: &Usage) {
let counters = turn_prompt_tokens();
counters
.with_label_values(&["input"])
.inc_by(usage.input_tokens);
counters
.with_label_values(&["cache_read"])
.inc_by(usage.cache_read_input_tokens);
counters
.with_label_values(&["cache_creation"])
.inc_by(usage.cache_creation_input_tokens);
}
pub(crate) fn force() {
for kind in ["input", "cache_read", "cache_creation"] {
turn_prompt_tokens().with_label_values(&[kind]);
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use polyc_llm::Usage;
use prometheus::{Encoder as _, TextEncoder};
use super::record_turn;
#[test]
fn record_turn_is_visible_in_a_registry_scrape() {
record_turn(&Usage {
input_tokens: 1000,
output_tokens: 42,
cache_read_input_tokens: 800,
cache_creation_input_tokens: 100,
});
let mut buf = Vec::new();
TextEncoder::new()
.encode(&prometheus::default_registry().gather(), &mut buf)
.expect("encode");
let text = String::from_utf8(buf).expect("utf8");
assert!(
text.contains("polychrome_turn_prompt_tokens_total"),
"missing counter in scrape:\n{text}"
);
assert!(text.contains("token_kind=\"input\""));
assert!(text.contains("token_kind=\"cache_read\""));
assert!(text.contains("token_kind=\"cache_creation\""));
}
#[test]
fn record_turn_with_no_cache_read_still_counts_input() {
record_turn(&Usage {
input_tokens: 500,
output_tokens: 10,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
});
let mut buf = Vec::new();
TextEncoder::new()
.encode(&prometheus::default_registry().gather(), &mut buf)
.expect("encode");
let text = String::from_utf8(buf).expect("utf8");
assert!(text.contains("polychrome_turn_prompt_tokens_total"));
}
}