rig_memory/
tokens.rs

1//! Estimating how many tokens a message, and the next request, cost.
2//!
3//! The next request is estimated from the last usage a provider reported
4//! while it still describes the request ([`next_request_tokens`]), and
5//! from the request as a whole otherwise.
6
7use std::hash::{DefaultHasher, Hash, Hasher};
8use std::io;
9use std::sync::Arc;
10
11use rig_core::completion::{Message, ToolDefinition, Usage};
12use rig_core::conversation::{Conversation, MessageId};
13use rig_core::wasm_compat::{WasmCompatSend, WasmCompatSync};
14use serde::{Deserialize, Serialize};
15
16/// Counts the tokens of one [`Message`]: the provider's tokenizer (such as
17/// `tiktoken-rs` for OpenAI) or an estimate. It runs once per message a
18/// window or a compaction weighs, so it must be cheap.
19///
20/// A counter chosen at runtime is shared as `Arc<dyn TokenCounter>`, owned as
21/// `Box<dyn TokenCounter>` or borrowed as `&dyn TokenCounter`; each counts as
22/// the counter it holds. (A generic `Box<T>` or `&T` would overlap the
23/// closure impl, since a boxed or borrowed closure is a closure too.)
24///
25/// ```
26/// use std::sync::Arc;
27/// use rig_core::completion::Message;
28/// use rig_memory::{HeuristicTokenCounter, TokenCounter, recent::TokenBudget};
29///
30/// let counter: Arc<dyn TokenCounter> = Arc::new(HeuristicTokenCounter::default());
31/// let budget = TokenBudget { max_tokens: 1_000, counter: Arc::clone(&counter) };
32/// assert_eq!(budget.counter.count(&Message::user("12345678")), 2 + 4);
33/// ```
34pub trait TokenCounter: WasmCompatSend + WasmCompatSync {
35    /// About how many tokens `message` costs.
36    fn count(&self, message: &Message) -> usize;
37}
38
39impl<F> TokenCounter for F
40where
41    F: Fn(&Message) -> usize + WasmCompatSend + WasmCompatSync,
42{
43    fn count(&self, message: &Message) -> usize {
44        (self)(message)
45    }
46}
47
48impl<C: TokenCounter + ?Sized> TokenCounter for Arc<C> {
49    fn count(&self, message: &Message) -> usize {
50        (**self).count(message)
51    }
52}
53
54impl TokenCounter for Box<dyn TokenCounter> {
55    fn count(&self, message: &Message) -> usize {
56        (**self).count(message)
57    }
58}
59
60impl TokenCounter for &dyn TokenCounter {
61    fn count(&self, message: &Message) -> usize {
62        (**self).count(message)
63    }
64}
65
66/// Estimates token cost from UTF-8 lengths, rounded up per text-bearing part.
67/// Tool calls count name bytes plus serialized argument bytes. Attachments and
68/// message overhead have fixed costs.
69///
70/// ```
71/// use rig_core::completion::Message;
72/// use rig_memory::{HeuristicTokenCounter, TokenCounter};
73///
74/// let counter = HeuristicTokenCounter::default();
75/// assert_eq!(counter.count(&Message::user("12345678")), 2 + 4);
76/// ```
77#[derive(Debug, Clone, Copy)]
78pub struct HeuristicTokenCounter {
79    bytes_per_token: f32,
80    per_message_overhead: usize,
81    per_attachment_tokens: usize,
82}
83
84impl HeuristicTokenCounter {
85    /// A counter with these parameters.
86    ///
87    /// Non-finite ratios and ratios below `1.0` become `1.0`.
88    /// Overhead and attachment costs may be zero.
89    pub fn new(
90        bytes_per_token: f32,
91        per_message_overhead: usize,
92        per_attachment_tokens: usize,
93    ) -> Self {
94        let bytes_per_token = if bytes_per_token.is_finite() && bytes_per_token >= 1.0 {
95            bytes_per_token
96        } else {
97            1.0
98        };
99        Self {
100            bytes_per_token,
101            per_message_overhead,
102            per_attachment_tokens,
103        }
104    }
105
106    /// Uses 3.5 bytes per token, four tokens per message, and
107    /// [`ATTACHMENT_TOKENS`] per attachment.
108    pub fn anthropic() -> Self {
109        Self::new(3.5, 4, ATTACHMENT_TOKENS)
110    }
111
112    fn bytes_to_tokens(&self, bytes: usize) -> usize {
113        // `bytes_per_token` is clamped to >= 1.0 in the constructor, so the
114        // division is well-defined. We round up so a single non-empty
115        // input still costs at least one token.
116        let tokens = (bytes as f32) / self.bytes_per_token;
117        tokens.ceil() as usize
118    }
119
120    /// About how many tokens the definition of a tool costs as a request
121    /// sends it: its name, description and parameters' JSON.
122    pub fn count_tool(&self, tool: &ToolDefinition) -> usize {
123        let bytes = tool.name.len().saturating_add(tool.description.len());
124        self.bytes_to_tokens(bytes.saturating_add(json_len(&tool.parameters)))
125            .saturating_add(self.per_message_overhead)
126    }
127
128    /// About how many tokens one item of a user message costs, without the
129    /// per-message overhead.
130    pub(crate) fn count_user(&self, content: &rig_core::message::UserContent) -> usize {
131        use rig_core::message::UserContent;
132        match content {
133            UserContent::Text(text) => self.bytes_to_tokens(text.text.len()),
134            UserContent::ToolResult(result) => result
135                .content
136                .iter()
137                .map(|c| match c {
138                    rig_core::message::ToolResultContent::Text(t) => {
139                        self.bytes_to_tokens(t.text.len())
140                    }
141                    rig_core::message::ToolResultContent::Json { value } => {
142                        self.bytes_to_tokens(json_len(value))
143                    }
144                    rig_core::message::ToolResultContent::Image(_) => self.per_attachment_tokens,
145                })
146                .fold(0, usize::saturating_add),
147            UserContent::Image(_)
148            | UserContent::Audio(_)
149            | UserContent::Video(_)
150            | UserContent::Document(_) => self.per_attachment_tokens,
151        }
152    }
153
154    /// About how many tokens one item of an assistant message costs, without
155    /// the per-message overhead.
156    fn count_assistant(&self, content: &rig_core::message::AssistantContent) -> usize {
157        use rig_core::message::AssistantContent;
158        match content {
159            AssistantContent::Text(text) => self.bytes_to_tokens(text.text.len()),
160            AssistantContent::Reasoning(reasoning) => self.bytes_to_tokens(reasoning.text.len()),
161            AssistantContent::ToolCall(call) => {
162                let name_bytes = call.function.name.len();
163                let args_bytes = json_len(&call.function.arguments);
164                self.bytes_to_tokens(name_bytes + args_bytes)
165            }
166            AssistantContent::Image(_) => self.per_attachment_tokens,
167            // A provider-only item is sent back to the model that made it.
168            AssistantContent::Opaque(opaque) => self.bytes_to_tokens(json_len(&opaque.item)),
169            other => self.bytes_to_tokens(json_len(other)),
170        }
171    }
172}
173
174/// The bytes of `value`'s JSON text, counted as it is written; 0 when it
175/// does not serialize.
176fn json_len(value: &impl serde::Serialize) -> usize {
177    struct Len(usize);
178    impl io::Write for Len {
179        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
180            self.0 += bytes.len();
181            Ok(bytes.len())
182        }
183
184        fn flush(&mut self) -> io::Result<()> {
185            Ok(())
186        }
187    }
188    let mut len = Len(0);
189    serde_json::to_writer(&mut len, value).map_or(0, |()| len.0)
190}
191
192/// About how many tokens an image, or another attachment, costs.
193pub const ATTACHMENT_TOKENS: usize = 1_200;
194
195impl Default for HeuristicTokenCounter {
196    /// Uses four bytes per token, four tokens per message, and
197    /// [`ATTACHMENT_TOKENS`] per attachment.
198    fn default() -> Self {
199        Self::new(4.0, 4, ATTACHMENT_TOKENS)
200    }
201}
202
203impl TokenCounter for HeuristicTokenCounter {
204    fn count(&self, message: &Message) -> usize {
205        let content_tokens: usize = match message {
206            Message::User { content } => content
207                .iter()
208                .map(|c| self.count_user(c))
209                .fold(0, usize::saturating_add),
210            Message::Assistant(turn) => turn
211                .content
212                .iter()
213                .map(|c| self.count_assistant(c))
214                .fold(0, usize::saturating_add),
215            Message::System { content } => self.bytes_to_tokens(content.len()),
216        };
217        content_tokens.saturating_add(self.per_message_overhead)
218    }
219}
220
221/// The usage a provider reported for an agent's last reply, and the last
222/// entry of the conversation it counts: the reply, or the request's last
223/// message when the reply added none. It anchors the estimate of the next
224/// request ([`next_request_tokens`]).
225#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
226pub struct Anchor {
227    /// What the provider reported. Its input counts every token the
228    /// request sent, the tool definitions included.
229    pub usage: Usage,
230    /// The last entry it counts.
231    pub through: MessageId,
232}
233
234/// A request as prepared to be sent, as [`next_request_tokens`] reads it.
235#[derive(Clone, Copy, Debug)]
236pub struct Prepared<'a> {
237    /// The model it goes to.
238    pub model: &'a str,
239    /// The system prompt sent before the messages, if any.
240    pub prompt: Option<&'a str>,
241    /// The tools it offers.
242    pub tools: &'a [Arc<ToolDefinition>],
243    /// Its messages, with their ids.
244    pub history: &'a Conversation,
245}
246
247/// What a sent request started with, to tell whether the usage reported
248/// for its reply still describes the next request: a fingerprint of its
249/// model, system prompt, tools and messages (each by id, the time it was
250/// made and the shared message it holds, so an edit, a splice or a restore
251/// changes it), and how many messages it had. It holds in the process that
252/// made it only.
253#[derive(Clone, Copy, Debug, PartialEq, Eq)]
254pub struct RequestStart {
255    fingerprint: u64,
256    messages: usize,
257}
258
259impl Prepared<'_> {
260    /// What this request starts with, to keep once it is sent.
261    pub fn start(&self) -> RequestStart {
262        let messages = self.history.messages().len();
263        RequestStart {
264            fingerprint: self.fingerprint(messages),
265            messages,
266        }
267    }
268
269    /// The fingerprint of the model, prompt, tools and first `messages`
270    /// messages.
271    fn fingerprint(&self, messages: usize) -> u64 {
272        let mut hasher = DefaultHasher::new();
273        (self.model, self.prompt).hash(&mut hasher);
274        for tool in self.tools {
275            let mut bytes = HashWriter(&mut hasher);
276            let _ = serde_json::to_writer(&mut bytes, &**tool);
277        }
278        for entry in self.history.entries().take(messages) {
279            (entry.key(), std::ptr::from_ref(entry.message).addr()).hash(&mut hasher);
280        }
281        hasher.finish()
282    }
283
284    /// The estimate of the whole request: its system prompt, tools and
285    /// messages, by the default [`HeuristicTokenCounter`].
286    pub fn whole_tokens(&self) -> u64 {
287        let counter = HeuristicTokenCounter::default();
288        let prompt = self
289            .prompt
290            .map_or(0, |prompt| counter.count(&Message::system(prompt)));
291        let tools = self.tools.iter().map(|tool| counter.count_tool(tool));
292        let tools = tools.fold(0, usize::saturating_add);
293        let messages = estimate_messages(self.history.messages());
294        (prompt.saturating_add(tools) as u64).saturating_add(messages)
295    }
296}
297
298/// Feeds bytes written to it to a hasher.
299struct HashWriter<'h>(&'h mut DefaultHasher);
300
301impl io::Write for HashWriter<'_> {
302    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
303        self.0.write(bytes);
304        Ok(bytes.len())
305    }
306
307    fn flush(&mut self) -> io::Result<()> {
308        Ok(())
309    }
310}
311
312/// The estimated tokens of `messages`, by the default
313/// [`HeuristicTokenCounter`].
314fn estimate_messages(messages: &[Arc<Message>]) -> u64 {
315    let counter = HeuristicTokenCounter::default();
316    messages
317        .iter()
318        .map(|message| counter.count(message) as u64)
319        .fold(0, u64::saturating_add)
320}
321
322/// About how many tokens `request` sends. Anchored, it is the usage the
323/// provider reported for the reply to the request that started as `sent`
324/// says, plus the estimate of the messages after the anchor's entry: the
325/// anchor holds while that entry is the reply to that request (or its last
326/// message) and `request` starts as that request did, with the same
327/// model, system prompt, tools and messages. Otherwise, such as after a
328/// compaction, a model change, a prompt or tool change, or a cleared copy,
329/// it is [`Prepared::whole_tokens`]. The tool definitions are counted
330/// once either way: reported input counts them already.
331pub fn next_request_tokens(
332    anchor: Option<&Anchor>,
333    sent: Option<RequestStart>,
334    request: &Prepared<'_>,
335) -> u64 {
336    let anchored = anchor.zip(sent).and_then(|(anchor, sent)| {
337        let at = request.history.index_of(anchor.through)?;
338        let follows = at == sent.messages || at.saturating_add(1) == sent.messages;
339        let same = follows && request.fingerprint(sent.messages) == sent.fingerprint;
340        let reported = anchor.usage.context_tokens().filter(|_| same)?;
341        let after = request.history.messages().get(at.saturating_add(1)..);
342        Some(reported.saturating_add(estimate_messages(after.unwrap_or_default())))
343    });
344    anchored.unwrap_or_else(|| request.whole_tokens())
345}
346
347/// Tokens a request leaves free below the window, besides the output it
348/// asks for: a margin for the estimate's error.
349pub const OUTPUT_MARGIN: u64 = 4_096;
350
351/// The most output a request of `estimate` tokens may ask of a model with
352/// a `window` of tokens and an output limit of `most`: `most`, or what the
353/// window leaves after the request and [`OUTPUT_MARGIN`] when that is
354/// less. `None` when neither is known; `Some(0)` when the window leaves no
355/// room, which calls for compacting first, not for a request with no room
356/// to answer.
357pub fn output_room(window: Option<u64>, most: Option<u64>, estimate: u64) -> Option<u64> {
358    let room = window.map(|window| window.saturating_sub(estimate.saturating_add(OUTPUT_MARGIN)));
359    match (room, most) {
360        (Some(room), Some(most)) => Some(room.min(most)),
361        (room, most) => room.or(most),
362    }
363}
364
365#[cfg(test)]
366mod tests;
rig_memory/tokens.rs

rig_memory/
tokens.rs
