Skip to main content

codetether_agent/session/helper/
token.rs

1use crate::provider::{ContentPart, Message, ToolDefinition};
2use crate::rlm::RlmChunker;
3
4/// Delegate to the canonical implementation in [`crate::provider::limits`].
5///
6/// Kept here as a re-export so existing call sites don't need to update
7/// their import paths. New code may import
8/// [`crate::provider::limits::context_window_for_model`] directly.
9pub use crate::provider::limits::context_window_for_model;
10
11pub fn session_completion_max_tokens() -> usize {
12    std::env::var("CODETETHER_SESSION_MAX_TOKENS")
13        .ok()
14        .and_then(|v| v.parse::<usize>().ok())
15        .filter(|v| *v > 0)
16        .unwrap_or(8192)
17}
18
19pub fn estimate_tokens_for_part(part: &ContentPart) -> usize {
20    match part {
21        ContentPart::Text { text } => RlmChunker::estimate_tokens(text),
22        ContentPart::ToolResult { content, .. } => RlmChunker::estimate_tokens(content),
23        ContentPart::ToolCall {
24            id,
25            name,
26            arguments,
27            thought_signature,
28            ..
29        } => {
30            let mut s = String::new();
31            s.push_str(id);
32            s.push(' ');
33            s.push_str(name);
34            s.push(' ');
35            s.push_str(arguments);
36            if let Some(sig) = thought_signature {
37                s.push(' ');
38                s.push_str(sig);
39            }
40            RlmChunker::estimate_tokens(&s)
41        }
42        ContentPart::Thinking { text } => RlmChunker::estimate_tokens(text),
43        ContentPart::Image { .. } => 2000,
44        ContentPart::File { path, mime_type } => {
45            let mut s = String::new();
46            s.push_str(path);
47            if let Some(mt) = mime_type {
48                s.push(' ');
49                s.push_str(mt);
50            }
51            RlmChunker::estimate_tokens(&s)
52        }
53    }
54}
55
56pub fn estimate_tokens_for_messages(messages: &[Message]) -> usize {
57    messages
58        .iter()
59        .map(|m| {
60            m.content
61                .iter()
62                .map(estimate_tokens_for_part)
63                .sum::<usize>()
64        })
65        .sum()
66}
67
68pub fn estimate_tokens_for_tools(tools: &[ToolDefinition]) -> usize {
69    serde_json::to_string(tools)
70        .ok()
71        .map(|s| RlmChunker::estimate_tokens(&s))
72        .unwrap_or(0)
73}
74
75pub fn estimate_request_tokens(
76    system_prompt: &str,
77    messages: &[Message],
78    tools: &[ToolDefinition],
79) -> usize {
80    RlmChunker::estimate_tokens(system_prompt)
81        + estimate_tokens_for_messages(messages)
82        + estimate_tokens_for_tools(tools)
83}