1use crate::types::{ModelMessage, Role};
2
3pub fn estimate_tokens(text: &str) -> usize {
4 (text.len() as f64 / 4.0).ceil() as usize
5}
6
7pub fn estimate_messages_tokens(messages: &[ModelMessage]) -> usize {
8 messages
9 .iter()
10 .map(|msg| estimate_tokens(&msg.content) + 4)
11 .sum()
12}
13
14pub fn trim_to_token_budget(messages: Vec<ModelMessage>, max_tokens: usize) -> Vec<ModelMessage> {
15 let system: Vec<_> = messages
16 .iter()
17 .filter(|m| m.role == Role::System)
18 .cloned()
19 .collect();
20 let rest: Vec<_> = messages
21 .iter()
22 .filter(|m| m.role != Role::System)
23 .cloned()
24 .collect();
25
26 let mut tokens = estimate_messages_tokens(&system);
27 let mut result = Vec::new();
28
29 for msg in rest.iter().rev() {
30 let t = estimate_tokens(&msg.content) + 4;
31 if tokens + t > max_tokens {
32 break;
33 }
34 result.insert(0, msg.clone());
35 tokens += t;
36 }
37
38 let mut final_messages = system;
39 final_messages.extend(result);
40 final_messages
41}