vv_agent/memory/
token_utils.rs

1use serde_json::Value;
2use std::path::Path;
3use std::sync::OnceLock;
4
5use crate::types::Message;
6
7pub fn resolve_model_token_limits(
8    settings: &Value,
9    backend: &str,
10    model: &str,
11) -> (Option<u64>, Option<u64>) {
12    let backend = backend.trim();
13    let model = model.trim();
14    if backend.is_empty() || model.is_empty() {
15        return (None, None);
16    }
17    let Ok(resolved) = crate::config::resolve_model_endpoint(settings, backend, model) else {
18        return (None, None);
19    };
20    (resolved.context_length, resolved.max_output_tokens)
21}
22
23pub fn resolve_model_token_limits_from_file(
24    path: impl AsRef<Path>,
25    backend: &str,
26    model: &str,
27) -> (Option<u64>, Option<u64>) {
28    let Ok(settings) = crate::config::load_llm_settings_from_file(path) else {
29        return (None, None);
30    };
31    resolve_model_token_limits(&settings, backend, model)
32}
33
34pub fn compute_compaction_threshold(
35    configured_threshold: u64,
36    model_context_window: u64,
37    reserved_output_tokens: u64,
38    autocompact_buffer_tokens: u64,
39) -> u64 {
40    let derived = if model_context_window > 0 {
41        model_context_window
42            .saturating_sub(reserved_output_tokens)
43            .saturating_sub(autocompact_buffer_tokens)
44    } else {
45        0
46    };
47    match (configured_threshold, derived) {
48        (configured, derived) if configured > 0 && derived > 0 => configured.min(derived),
49        (configured, _) if configured > 0 => configured,
50        (_, derived) => derived,
51    }
52}
53
54pub fn count_messages_tokens(messages: &[Message], model: &str) -> u64 {
55    if messages.is_empty() {
56        return 0;
57    }
58    let combined_text = messages
59        .iter()
60        .map(|message| message.content.as_str())
61        .collect::<Vec<_>>()
62        .join("\n");
63    let text_tokens = if combined_text.is_empty() {
64        0
65    } else {
66        count_tokens(&combined_text, model)
67    };
68    let image_count = messages
69        .iter()
70        .filter(|message| message.image_url.is_some())
71        .count() as u64;
72
73    text_tokens + image_count * default_image_tokens(model)
74}
75
76pub trait TokenCountPayload {
77    fn to_token_count_text(&self) -> String;
78}
79
80impl TokenCountPayload for str {
81    fn to_token_count_text(&self) -> String {
82        self.to_string()
83    }
84}
85
86impl TokenCountPayload for String {
87    fn to_token_count_text(&self) -> String {
88        self.clone()
89    }
90}
91
92impl TokenCountPayload for Value {
93    fn to_token_count_text(&self) -> String {
94        serde_json::to_string(self).unwrap_or_default()
95    }
96}
97
98pub fn count_tokens<T: TokenCountPayload + ?Sized>(payload: &T, model: &str) -> u64 {
99    let text = payload.to_token_count_text();
100    if text.is_empty() {
101        return 0;
102    }
103    let direct_count = vv_llm::utilities::count_tokens(&text, model);
104    if vv_llm_has_universal_model_dispatch() {
105        if let Ok(count) = direct_count {
106            if count > 0 {
107                return count as u64;
108            }
109        }
110    } else if let Some(count) = count_tokens_with_legacy_vv_llm(&text, model) {
111        if count > 0 {
112            return count as u64;
113        }
114    }
115    estimate_tokens(&text, model)
116}
117
118pub fn estimate_tokens(text: &str, _model: &str) -> u64 {
119    if text.is_empty() {
120        return 0;
121    }
122    let mut cjk_chars = 0_u64;
123    let mut other_chars = 0_u64;
124    for ch in text.chars() {
125        if is_cjk(ch) {
126            cjk_chars += 1;
127        } else {
128            other_chars += 1;
129        }
130    }
131    ((cjk_chars as f64 * 1.5) + (other_chars as f64 * 0.25))
132        .floor()
133        .max(1.0) as u64
134}
135
136fn default_image_tokens(model: &str) -> u64 {
137    if model.to_ascii_lowercase().starts_with("moonshot") {
138        1_024
139    } else {
140        765
141    }
142}
143
144fn vv_llm_has_universal_model_dispatch() -> bool {
145    static SUPPORTED: OnceLock<bool> = OnceLock::new();
146    *SUPPORTED.get_or_init(|| {
147        matches!(
148            vv_llm::utilities::count_tokens(
149                "antidisestablishmentarianism",
150                "unknown-provider-model",
151            ),
152            Ok(6)
153        )
154    })
155}
156
157fn count_tokens_with_legacy_vv_llm(text: &str, model: &str) -> Option<usize> {
158    let normalized_model = model.to_ascii_lowercase();
159    if normalized_model.starts_with("abab") || normalized_model.starts_with("minimax") {
160        return Some((text.chars().count() as f64 / 1.33) as usize);
161    }
162
163    // vv-llm 0.2.x exposes both BPEs but only dispatches a few model names.
164    let tokenizer_model = if normalized_model == "gpt-3.5-turbo"
165        || normalized_model.starts_with("moonshot")
166        || normalized_model.starts_with("kimi")
167        || normalized_model.starts_with("gemini")
168        || normalized_model.starts_with("stepfun")
169        || normalized_model.starts_with("glm")
170    {
171        "gpt-3.5-turbo"
172    } else {
173        "gpt-4o"
174    };
175    vv_llm::utilities::count_tokens(text, tokenizer_model).ok()
176}
177
178fn is_cjk(ch: char) -> bool {
179    matches!(
180        ch as u32,
181        0x4E00..=0x9FFF | 0x3000..=0x303F | 0xFF00..=0xFFEF
182    )
183}
vv_agent/memory/token_utils.rs

vv_agent/memory/
token_utils.rs
