phi_core/context/
tracker.rs

1use super::token::{message_tokens, total_tokens};
2use crate::types::*;
3
4// ---------------------------------------------------------------------------
5// Context tracking (real usage + estimates)
6// ---------------------------------------------------------------------------
7
8/// Tracks context size using real token counts from provider responses
9/// combined with estimates for messages added after the last response.
10///
11/// This gives more accurate context size tracking than pure estimation,
12/// since providers report actual token counts in their usage data.
13///
14/// # Example
15///
16/// ```rust
17/// use phi_core::context::ContextTracker;
18/// use phi_core::types::Usage;
19///
20/// let mut tracker = ContextTracker::new();
21/// // After receiving an assistant response with usage data:
22/// tracker.record_usage(&Usage { input: 1500, output: 200, ..Default::default() }, 3);
23/// ```
24/*
25RUST QUIRK: Using `Option<usize>` for "not yet known" state
26
27`last_usage_tokens: Option<usize>` means "either we have a real token count
28(Some(n)), or we haven't received one yet (None)".
29
30This is Rust's way of representing nullable data without null pointers.
31There is no `null` or `None` in Rust — you must use `Option<T>` explicitly.
32The compiler forces you to handle both cases, preventing null pointer exceptions.
33
34Python analogy: last_usage_tokens: Optional[int] = None
35
36The hybrid design strategy:
37  - After each LLM response, record the REAL token count from provider usage data
38  - For messages added after the last response, ESTIMATE with chars/4
39  - Combine: real_base + estimated_trailing = accurate context size
40
41This beats pure estimation because real token counts account for:
42  - Unicode characters (multi-byte)
43  - Special tokens (BOS, EOS, system prompt formatting)
44  - Provider-specific tokenization differences
45*/
46pub struct ContextTracker {
47    /// Last known total token count from provider usage (None = no response yet)
48    last_usage_tokens: Option<usize>,
49    /// Index into the message list of the last assistant response with usage (None = no response yet)
50    last_usage_index: Option<usize>,
51}
52
53impl ContextTracker {
54    pub fn new() -> Self {
55        Self {
56            last_usage_tokens: None,
57            last_usage_index: None,
58        }
59    }
60
61    /// Record usage from an assistant response.
62    ///
63    /// Call this after each assistant message to update the tracker
64    /// with real token counts from the provider.
65    pub fn record_usage(&mut self, usage: &Usage, message_index: usize) {
66        let total = usage.input + usage.output + usage.cache_read + usage.cache_write;
67        if total > 0 {
68            self.last_usage_tokens = Some(total as usize);
69            self.last_usage_index = Some(message_index);
70        }
71    }
72
73    /// Estimate current context size.
74    ///
75    /// Uses real usage from the last assistant response as a baseline,
76    /// then adds estimates (chars/4) for any messages added since.
77    /// Falls back to pure estimation if no usage data is available.
78    pub fn estimate_context_tokens(&self, messages: &[AgentMessage]) -> usize {
79        match (self.last_usage_tokens, self.last_usage_index) {
80            (Some(usage_tokens), Some(idx)) if idx < messages.len() => {
81                let trailing: usize = messages[idx + 1..].iter().map(message_tokens).sum();
82                usage_tokens + trailing
83            }
84            _ => total_tokens(messages),
85        }
86    }
87
88    /// Reset tracking (e.g. after compaction replaces messages).
89    pub fn reset(&mut self) {
90        self.last_usage_tokens = None;
91        self.last_usage_index = None;
92    }
93}
94
95impl Default for ContextTracker {
96    fn default() -> Self {
97        Self::new()
98    }
99}
phi_core/context/tracker.rs

phi_core/context/
tracker.rs
