albatross-cli 2.0.3

Open source, terminal-first AI coding agent with fully transparent multi-model routing. Local (Ollama, LM Studio, MLX, llama.cpp) or cloud, your keys, one TUI.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use serde::Serialize;

use crate::openai::{ChatMessage, ToolDef};

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PromptBudget {
    pub system_bytes: usize,
    pub transcript_bytes: usize,
    pub tool_schema_bytes: usize,
    pub tool_result_bytes: usize,
    /// Bytes sent on the wire: system prompt + transcript JSON + tool schemas.
    pub total_bytes: usize,
    pub effective_total_bytes: usize,
    pub estimated_tokens: usize,
}

pub fn estimate_tokens(bytes: usize) -> usize {
    bytes.div_ceil(4)
}

pub fn measure_prompt_budget(
    system_prompt: &str,
    messages: &[ChatMessage],
    tool_defs: &[ToolDef],
) -> PromptBudget {
    let system_bytes = system_prompt.len();
    let transcript_bytes = messages
        .iter()
        .filter(|m| !matches!(m, ChatMessage::System { .. }))
        .map(|m| serde_json::to_vec(m).map(|v| v.len()).unwrap_or(0))
        .sum();
    let tool_result_bytes = messages
        .iter()
        .map(|m| match m {
            ChatMessage::Tool { content, .. } => content.len(),
            _ => 0,
        })
        .sum();
    let tool_schema_bytes = if tool_defs.is_empty() {
        0
    } else {
        serde_json::to_vec(tool_defs).map(|v| v.len()).unwrap_or(0)
    };
    let total_bytes = system_bytes + transcript_bytes + tool_schema_bytes;
    PromptBudget {
        system_bytes,
        transcript_bytes,
        tool_schema_bytes,
        tool_result_bytes,
        effective_total_bytes: total_bytes,
        total_bytes,
        estimated_tokens: estimate_tokens(total_bytes),
    }
}

pub fn usage_ratio(budget: &PromptBudget, limit_bytes: usize) -> f64 {
    if limit_bytes == 0 {
        return 0.0;
    }
    budget.effective_total_bytes as f64 / limit_bytes as f64
}

pub fn headroom_bytes(budget: &PromptBudget, limit_bytes: usize) -> usize {
    limit_bytes.saturating_sub(budget.effective_total_bytes)
}

pub fn format_bytes(n: usize) -> String {
    if n >= 1024 * 1024 {
        format!("{:.1} MB", n as f64 / (1024.0 * 1024.0))
    } else if n >= 1024 {
        format!("{:.1} KB", n as f64 / 1024.0)
    } else {
        format!("{n} B")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn estimates_tokens_at_four_bytes_each() {
        assert_eq!(estimate_tokens(0), 0);
        assert_eq!(estimate_tokens(1), 1);
        assert_eq!(estimate_tokens(8), 2);
    }

    #[test]
    fn separates_tool_result_bytes() {
        let messages = vec![ChatMessage::Tool {
            tool_call_id: "x".into(),
            content: "hello".into(),
        }];
        let budget = measure_prompt_budget("system", &messages, &[]);
        assert_eq!(budget.system_bytes, 6);
        assert_eq!(budget.tool_result_bytes, 5);
        assert!(budget.transcript_bytes >= 5);
        assert_eq!(budget.total_bytes, budget.effective_total_bytes);
        assert!(budget.effective_total_bytes >= budget.tool_result_bytes);
    }

    #[test]
    fn usage_ratio_and_headroom() {
        let budget = PromptBudget {
            system_bytes: 100,
            transcript_bytes: 300,
            tool_schema_bytes: 100,
            tool_result_bytes: 50,
            total_bytes: 500,
            effective_total_bytes: 500,
            estimated_tokens: 125,
        };
        assert!((usage_ratio(&budget, 1000) - 0.5).abs() < f64::EPSILON);
        assert_eq!(headroom_bytes(&budget, 1000), 500);
    }
}