nexus_memory_agent/
token_budget.rs

1//! Token budget management for context window allocation.
2
3/// Token budget allocation for a model's context window.
4#[derive(Debug, Clone, Copy)]
5pub struct TokenBudget {
6    /// Total context window of the model
7    pub model_context_window: usize,
8    /// Percentage of window allocated to Nexus (0.0 - 1.0)
9    pub nexus_allocation_pct: f32,
10    /// Budget for the soul.md document
11    pub soul_budget: usize,
12    /// Budget for the project context.md (hot cache + cold recall)
13    pub context_budget: usize,
14}
15
16impl TokenBudget {
17    /// Create a token budget based on a model's context window.
18    pub fn for_model(context_window: usize) -> Self {
19        // Default allocation: 10% of window to Nexus
20        let nexus_pct = 0.10;
21        let nexus_budget = (context_window as f32 * nexus_pct) as usize;
22
23        // Of the Nexus budget: 30% for soul, 70% for context
24        let soul_budget = (nexus_budget as f32 * 0.30) as usize;
25        let context_budget = nexus_budget - soul_budget;
26
27        Self {
28            model_context_window: context_window,
29            nexus_allocation_pct: nexus_pct,
30            soul_budget,
31            context_budget,
32        }
33    }
34
35    /// Estimate context window for known agent types/models.
36    pub fn estimate_window(agent_type: &str) -> usize {
37        match agent_type.to_lowercase().as_str() {
38            "claude-code" | "amp" | "codex" | "sonnet-specialist" => 200_000,
39            "gemini" | "gemini-analyzer" => 1_000_000,
40            "opus-specialist" => 200_000,
41            _ => 128_000, // Conservative default for GPT-4/Turbo/o1
42        }
43    }
44
45    /// Rough estimation of tokens from text length (approx 4 chars per token).
46    pub fn estimate_tokens(text: &str) -> usize {
47        text.len() / 4
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn test_token_budget_math() {
57        let budget = TokenBudget::for_model(100_000);
58        assert_eq!(budget.model_context_window, 100_000);
59        // 10% of 100k = 10k
60        // soul = 30% of 10k = 3k
61        // context = 7k
62        assert_eq!(budget.soul_budget, 3000);
63        assert_eq!(budget.context_budget, 7000);
64    }
65
66    #[test]
67    fn test_estimate_window() {
68        assert_eq!(TokenBudget::estimate_window("claude-code"), 200_000);
69        assert_eq!(TokenBudget::estimate_window("gemini"), 1_000_000);
70        assert_eq!(TokenBudget::estimate_window("unknown"), 128_000);
71    }
72}
nexus_memory_agent/token_budget.rs

nexus_memory_agent/
token_budget.rs
