vv_agent/memory/
token_utils.rs

1use serde_json::Value;
2use std::path::Path;
3
4use crate::types::Message;
5
6pub fn resolve_model_token_limits(
7    settings: &Value,
8    backend: &str,
9    model: &str,
10) -> (Option<u64>, Option<u64>) {
11    let backend = backend.trim();
12    let model = model.trim();
13    if backend.is_empty() || model.is_empty() {
14        return (None, None);
15    }
16    let Ok(resolved) = crate::config::resolve_model_endpoint(settings, backend, model) else {
17        return (None, None);
18    };
19    (resolved.context_length, resolved.max_output_tokens)
20}
21
22pub fn resolve_model_token_limits_from_file(
23    path: impl AsRef<Path>,
24    backend: &str,
25    model: &str,
26) -> (Option<u64>, Option<u64>) {
27    let Ok(settings) = crate::config::load_llm_settings_from_file(path) else {
28        return (None, None);
29    };
30    resolve_model_token_limits(&settings, backend, model)
31}
32
33pub fn compute_compaction_threshold(
34    configured_threshold: u64,
35    model_context_window: u64,
36    reserved_output_tokens: u64,
37    autocompact_buffer_tokens: u64,
38) -> u64 {
39    let derived = if model_context_window > 0 {
40        model_context_window
41            .saturating_sub(reserved_output_tokens)
42            .saturating_sub(autocompact_buffer_tokens)
43    } else {
44        0
45    };
46    if configured_threshold > 0 {
47        configured_threshold.min(derived)
48    } else {
49        derived
50    }
51}
52
53pub fn count_messages_tokens(messages: &[Message], model: &str) -> u64 {
54    if messages.is_empty() {
55        return 0;
56    }
57    let combined_text = messages
58        .iter()
59        .map(|message| message.content.as_str())
60        .collect::<Vec<_>>()
61        .join("\n");
62    let text_tokens = if combined_text.is_empty() {
63        0
64    } else {
65        count_tokens(&combined_text, model)
66    };
67    let image_count = messages
68        .iter()
69        .filter(|message| message.image_url.is_some())
70        .count() as u64;
71
72    text_tokens + image_count * default_image_tokens(model)
73}
74
75pub trait TokenCountPayload {
76    fn to_token_count_text(&self) -> String;
77}
78
79impl TokenCountPayload for str {
80    fn to_token_count_text(&self) -> String {
81        self.to_string()
82    }
83}
84
85impl TokenCountPayload for String {
86    fn to_token_count_text(&self) -> String {
87        self.clone()
88    }
89}
90
91impl TokenCountPayload for Value {
92    fn to_token_count_text(&self) -> String {
93        serde_json::to_string(self).unwrap_or_default()
94    }
95}
96
97pub fn count_tokens<T: TokenCountPayload + ?Sized>(payload: &T, model: &str) -> u64 {
98    let text = payload.to_token_count_text();
99    if text.is_empty() {
100        return 0;
101    }
102    if let Ok(count) = vv_llm::utilities::count_tokens(&text, model) {
103        if count > 0 {
104            return count as u64;
105        }
106    }
107    estimate_tokens(&text, model)
108}
109
110pub fn estimate_tokens(text: &str, _model: &str) -> u64 {
111    if text.is_empty() {
112        return 0;
113    }
114    let mut cjk_chars = 0_u64;
115    let mut other_chars = 0_u64;
116    for ch in text.chars() {
117        if is_cjk(ch) {
118            cjk_chars += 1;
119        } else {
120            other_chars += 1;
121        }
122    }
123    ((cjk_chars as f64 * 1.5) + (other_chars as f64 * 0.25))
124        .floor()
125        .max(1.0) as u64
126}
127
128fn default_image_tokens(model: &str) -> u64 {
129    if model.to_ascii_lowercase().starts_with("moonshot") {
130        1_024
131    } else {
132        765
133    }
134}
135
136fn is_cjk(ch: char) -> bool {
137    matches!(
138        ch as u32,
139        0x4E00..=0x9FFF | 0x3000..=0x303F | 0xFF00..=0xFFEF
140    )
141}
vv_agent/memory/token_utils.rs

vv_agent/memory/
token_utils.rs
