codetether_agent/session/helper/
token.rs1use crate::provider::{ContentPart, Message, ToolDefinition};
2use crate::rlm::RlmChunker;
3
4pub 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}