Skip to main content

nexus_memory_agent/
context_builder.rs

1//! Builds context.md from hot and cold caches with budget awareness.
2
3use crate::cognitive_cache::{ConfidenceTier, HotCache, HotCacheEntry};
4use crate::token_budget::TokenBudget;
5
6/// Compression level for a memory entry.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum CompressionLevel {
9    /// Full content
10    None,
11    /// First sentence or brief summary
12    Light,
13    /// One-liner
14    Heavy,
15}
16
17/// A recalled memory from the cold index or vector search.
18#[derive(Debug, Clone)]
19pub struct ColdRecall {
20    pub memory_id: i64,
21    pub content: String,
22    pub relevance_score: f32,
23    pub tier: ConfidenceTier,
24}
25
26/// Builds the context.md string from hot and cold memories.
27pub fn build_context_md(hot: &HotCache, cold: &[ColdRecall], max_tokens: usize) -> String {
28    if max_tokens == 0 {
29        return String::new();
30    }
31
32    let mut output = String::new();
33    let mut current_tokens = 0;
34
35    output.push_str("# Nexus Project Context\n\n");
36
37    // Group hot entries by tier
38    let loud_entries: Vec<&HotCacheEntry> = hot
39        .entries
40        .iter()
41        .filter(|e| e.tier == ConfidenceTier::Loud)
42        .collect();
43    let clear_entries: Vec<&HotCacheEntry> = hot
44        .entries
45        .iter()
46        .filter(|e| e.tier == ConfidenceTier::Clear)
47        .collect();
48    let whisper_entries: Vec<&HotCacheEntry> = hot
49        .entries
50        .iter()
51        .filter(|e| e.tier == ConfidenceTier::Whisper)
52        .collect();
53
54    // 1. Add Loud entries (Full content)
55    let loud_section = format_tier_section(
56        "## High Relevance (Loud)",
57        &loud_entries,
58        CompressionLevel::None,
59    );
60    let loud_tokens = TokenBudget::estimate_tokens(&loud_section);
61    if current_tokens + loud_tokens <= max_tokens {
62        output.push_str(&loud_section);
63        current_tokens += loud_tokens;
64    } else {
65        // Truncation happens if even Loud doesn't fit
66        return output;
67    }
68
69    // 2. Add Clear entries (Light compression)
70    let clear_section = format_tier_section(
71        "## Relevant (Clear)",
72        &clear_entries,
73        CompressionLevel::Light,
74    );
75    let clear_tokens = TokenBudget::estimate_tokens(&clear_section);
76    if current_tokens + clear_tokens <= max_tokens {
77        output.push_str(&clear_section);
78        current_tokens += clear_tokens;
79    }
80
81    // 3. Add Whisper entries (Heavy compression) - only if < 80% budget used
82    if (current_tokens as f32 / max_tokens as f32) < 0.80 {
83        let whisper_section = format_tier_section(
84            "## Low Signal (Whisper)",
85            &whisper_entries,
86            CompressionLevel::Heavy,
87        );
88        let whisper_tokens = TokenBudget::estimate_tokens(&whisper_section);
89        if current_tokens + whisper_tokens <= max_tokens {
90            output.push_str(&whisper_section);
91            current_tokens += whisper_tokens;
92        }
93    }
94
95    // 4. Add Cold Recalls
96    if !cold.is_empty() && max_tokens - current_tokens > 200 {
97        output.push_str("## Recalled Memories\n\n");
98        for recall in cold {
99            let entry_str = format!(
100                "- [Recall {}] {}\n",
101                recall.memory_id,
102                compress_text(&recall.content, CompressionLevel::Light)
103            );
104            let entry_tokens = TokenBudget::estimate_tokens(&entry_str);
105            if current_tokens + entry_tokens <= max_tokens {
106                output.push_str(&entry_str);
107                current_tokens += entry_tokens;
108            } else {
109                break;
110            }
111        }
112    }
113
114    output
115}
116
117fn format_tier_section(
118    title: &str,
119    entries: &[&HotCacheEntry],
120    compression: CompressionLevel,
121) -> String {
122    if entries.is_empty() {
123        return String::new();
124    }
125
126    let mut section = format!("{}\n\n", title);
127    for entry in entries {
128        let content = compress_text(&entry.content, compression);
129        section.push_str(&format!("### Memory {}\n{}\n\n", entry.memory_id, content));
130    }
131    section
132}
133
134fn compress_text(text: &str, level: CompressionLevel) -> String {
135    match level {
136        CompressionLevel::None => text.to_string(),
137        CompressionLevel::Light => text.lines().next().unwrap_or("").to_string(),
138        CompressionLevel::Heavy => {
139            let first_line = text.lines().next().unwrap_or("");
140            if first_line.len() > 80 {
141                // UTF-8-safe: find the char boundary at or before byte 77
142                let truncate_at = first_line
143                    .char_indices()
144                    .take_while(|(idx, _)| *idx < 77)
145                    .last()
146                    .map(|(idx, c)| idx + c.len_utf8())
147                    .unwrap_or(0);
148                format!("{}...", &first_line[..truncate_at])
149            } else {
150                first_line.to_string()
151            }
152        }
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use crate::cognitive_cache::ConfidenceTier;
160    use chrono::Utc;
161
162    #[test]
163    fn test_build_context_md_ordering_and_budget() {
164        let mut hot = HotCache::default();
165        hot.entries.push(HotCacheEntry {
166            memory_id: 1,
167            content: "Loud Content".into(),
168            relevance_score: 0.9,
169            tier: ConfidenceTier::Loud,
170            promoted_at: Utc::now(),
171            last_surfaced: Utc::now(),
172            hot_streak: 1,
173            pinned: false,
174            source_agent: None,
175        });
176        hot.entries.push(HotCacheEntry {
177            memory_id: 2,
178            content: "Clear Content Line 1\nLine 2".into(),
179            relevance_score: 0.75,
180            tier: ConfidenceTier::Clear,
181            promoted_at: Utc::now(),
182            last_surfaced: Utc::now(),
183            hot_streak: 1,
184            pinned: false,
185            source_agent: None,
186        });
187
188        // 1000 tokens budget (plenty)
189        let context = build_context_md(&hot, &[], 1000);
190        assert!(context.contains("High Relevance (Loud)"));
191        assert!(context.contains("Loud Content"));
192        assert!(context.contains("Relevant (Clear)"));
193        assert!(context.contains("Clear Content Line 1"));
194        assert!(!context.contains("Line 2")); // Light compression
195
196        // Very small budget
197        let tight_context = build_context_md(&hot, &[], 10);
198        assert!(tight_context.len() < context.len());
199    }
200
201    #[test]
202    fn test_build_context_md_with_cold_recall() {
203        let hot = HotCache::default();
204        let cold = vec![ColdRecall {
205            memory_id: 99,
206            content: "Cold Content".into(),
207            relevance_score: 0.68,
208            tier: ConfidenceTier::Whisper,
209        }];
210
211        let context = build_context_md(&hot, &cold, 1000);
212        assert!(context.contains("Recalled Memories"));
213        assert!(context.contains("[Recall 99] Cold Content"));
214    }
215
216    #[test]
217    fn test_compress_text_utf8_multibyte() {
218        // Japanese text: each char is 3 bytes. 80+ byte line triggers Heavy truncation.
219        let long_multibyte = "ア".repeat(40); // 120 bytes, 40 chars
220        let compressed = compress_text(&long_multibyte, CompressionLevel::Heavy);
221        assert!(compressed.ends_with("..."));
222        // Must not panic and must be valid UTF-8
223        assert!(compressed.is_char_boundary(compressed.len()));
224        // Should be under 80 bytes + "..."
225        assert!(compressed.len() <= 83);
226    }
227}