htmlp/runtime/
usage.rs

1use super::{ContentBlock, ContentSource, ModelRequest};
2use serde::{Deserialize, Serialize};
3use std::{collections::BTreeMap, num::NonZeroUsize};
4
5/// Explicit heuristic policy. Image transport bytes are never treated as text.
6/// This estimate excludes provider framing and is not a tokenizer count.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
8#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
9pub struct ByteEstimator {
10    pub bytes_per_token: NonZeroUsize,
11    pub image_tokens: u64,
12}
13impl Default for ByteEstimator {
14    fn default() -> Self {
15        Self {
16            bytes_per_token: NonZeroUsize::new(4).unwrap(),
17            image_tokens: 8192,
18        }
19    }
20}
21impl ByteEstimator {
22    /// Round up so nonempty text never disappears from accounting.
23    pub fn text_tokens(&self, text: &str) -> u64 {
24        self.byte_tokens(text.len())
25    }
26    fn byte_tokens(&self, bytes: usize) -> u64 {
27        bytes.div_ceil(self.bytes_per_token.get()) as u64
28    }
29    /// Aggregate bytes by the full typed source before rounding each source
30    /// once. The total is exactly the sum of these non-overlapping categories.
31    pub fn estimate(&self, request: &ModelRequest) -> RequestEstimate {
32        let mut sources: BTreeMap<ContentSource, (usize, u64)> = BTreeMap::new();
33        for fragment in request.fragments() {
34            let (bytes, images) = sources.entry(fragment.source().clone()).or_default();
35            for block in fragment.content() {
36                match block {
37                    ContentBlock::Image { .. } => *images = images.saturating_add(1),
38                    ContentBlock::Text { .. }
39                    | ContentBlock::Reasoning { .. }
40                    | ContentBlock::ToolUse { .. }
41                    | ContentBlock::ToolResult { .. } => {
42                        *bytes = bytes.saturating_add(block.byte_len());
43                    }
44                }
45            }
46        }
47        if !request.tools().is_empty() {
48            let (bytes, _) = sources.entry(ContentSource::ToolDefinitions).or_default();
49            for tool in request.tools() {
50                *bytes = bytes.saturating_add(tool.byte_len());
51            }
52        }
53        let sources = sources
54            .into_iter()
55            .map(|(source, (text_bytes, images))| SourceEstimate {
56                source,
57                text_bytes,
58                images,
59                tokens: self
60                    .byte_tokens(text_bytes)
61                    .saturating_add(images.saturating_mul(self.image_tokens)),
62            })
63            .collect();
64        RequestEstimate {
65            method: *self,
66            sources,
67        }
68    }
69}
70
71/// One source's heuristic allocation, not provider-reported category usage.
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
74pub struct SourceEstimate {
75    pub source: ContentSource,
76    pub text_bytes: usize,
77    pub images: u64,
78    pub tokens: u64,
79}
80
81/// Categorized estimate. Total derives from categories instead of being a
82/// separately mutable field that could drift from them.
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
85pub struct RequestEstimate {
86    method: ByteEstimator,
87    sources: Vec<SourceEstimate>,
88}
89impl RequestEstimate {
90    pub fn method(&self) -> ByteEstimator {
91        self.method
92    }
93    pub fn sources(&self) -> &[SourceEstimate] {
94        &self.sources
95    }
96    pub fn total_tokens(&self) -> u64 {
97        self.sources
98            .iter()
99            .fold(0u64, |total, item| total.saturating_add(item.tokens))
100    }
101    pub fn tokens_for(&self, source: &ContentSource) -> u64 {
102        self.sources
103            .iter()
104            .find(|s| &s.source == source)
105            .map_or(0, |s| s.tokens)
106    }
107}
108
109/// Provider-reported usage for one model call. Input/cache lanes are disjoint;
110/// reasoning is a subset of output and must not be added to the output total.
111#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
112#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
113pub struct ReportedUsage {
114    pub input_tokens: u64,
115    pub output_tokens: u64,
116    pub cache_read_tokens: u64,
117    pub cache_write_tokens: u64,
118    pub reasoning_tokens: u64,
119}
120impl ReportedUsage {
121    pub fn prompt_tokens(&self) -> u64 {
122        self.input_tokens
123            .saturating_add(self.cache_read_tokens)
124            .saturating_add(self.cache_write_tokens)
125    }
126    /// Merge cumulative, partially populated snapshots of the SAME call.
127    /// Do not use this to aggregate distinct requests.
128    pub fn merge_snapshot(&mut self, snapshot: &Self) {
129        self.input_tokens = self.input_tokens.max(snapshot.input_tokens);
130        self.output_tokens = self.output_tokens.max(snapshot.output_tokens);
131        self.cache_read_tokens = self.cache_read_tokens.max(snapshot.cache_read_tokens);
132        self.cache_write_tokens = self.cache_write_tokens.max(snapshot.cache_write_tokens);
133        self.reasoning_tokens = self.reasoning_tokens.max(snapshot.reasoning_tokens);
134    }
135    /// Sum distinct calls. Reasoning stays an explanatory subset of output.
136    pub fn add_call(&mut self, call: &Self) {
137        self.input_tokens = self.input_tokens.saturating_add(call.input_tokens);
138        self.output_tokens = self.output_tokens.saturating_add(call.output_tokens);
139        self.cache_read_tokens = self
140            .cache_read_tokens
141            .saturating_add(call.cache_read_tokens);
142        self.cache_write_tokens = self
143            .cache_write_tokens
144            .saturating_add(call.cache_write_tokens);
145        self.reasoning_tokens = self.reasoning_tokens.saturating_add(call.reasoning_tokens);
146    }
147}
148
149/// Usage for one call. `estimate: None` means source attribution is unavailable
150/// (for example, a backend assembled the prompt), not zero category usage.
151#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
152#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
153pub struct TokenUsage {
154    pub estimate: Option<RequestEstimate>,
155    pub reported: Option<ReportedUsage>,
156    pub context_window_tokens: u64,
157}
158impl TokenUsage {
159    /// Best available prompt fill: reported input including cache lanes,
160    /// otherwise the categorized estimate. None means no measurement exists.
161    pub fn prompt_tokens(&self) -> Option<u64> {
162        self.reported
163            .as_ref()
164            .map(ReportedUsage::prompt_tokens)
165            .or_else(|| self.estimate.as_ref().map(RequestEstimate::total_tokens))
166    }
167}
htmlp/runtime/usage.rs

htmlp/runtime/
usage.rs
