Skip to main content

tandem_memory/
context_layers.rs

1use crate::types::{LayerType, MemoryError, MemoryResult};
2use std::sync::Arc;
3use tandem_providers::ProviderRegistry;
4
5const L0_SUMMARY_PROMPT: &str = r#"You are a text abstraction engine. Given the input text, generate a very brief abstract (~100 tokens or less) that captures the core essence.
6
7Requirements:
8- One sentence that summarizes the entire content
9- Focus on what this is about, not details
10- Be concise and direct
11- Start with the main topic/action
12
13Input text:
14"#;
15
16const L1_OVERVIEW_PROMPT: &str = r#"You are a text summarization engine. Given the input text, generate an overview (~2000 tokens or less) that captures the key information.
17
18Requirements:
19- Start with a brief introduction (1-2 sentences)
20- Cover main topics, key decisions, important facts
21- Omit conversational filler, greetings, and sign-offs
22- Use bullet points for lists where appropriate
23- End with any unresolved issues or follow-up items if relevant
24
25Input text:
26"#;
27
28pub struct ContextLayerGenerator {
29    providers: Arc<ProviderRegistry>,
30}
31
32impl ContextLayerGenerator {
33    pub fn new(providers: Arc<ProviderRegistry>) -> Self {
34        Self { providers }
35    }
36
37    pub async fn generate_l0(&self, content: &str) -> MemoryResult<String> {
38        let prompt = format!("{}{}", L0_SUMMARY_PROMPT, content);
39        self.call_llm(&prompt, "summarize-l0").await
40    }
41
42    pub async fn generate_l1(&self, content: &str) -> MemoryResult<String> {
43        let prompt = format!("{}{}", L1_OVERVIEW_PROMPT, content);
44        self.call_llm(&prompt, "summarize-l1").await
45    }
46
47    pub async fn generate_layers(&self, l2_content: &str) -> MemoryResult<(String, String)> {
48        let l0 = self.generate_l0(l2_content).await?;
49        let l1 = self.generate_l1(l2_content).await?;
50        Ok((l0, l1))
51    }
52
53    async fn call_llm(&self, prompt: &str, _model_hint: &str) -> MemoryResult<String> {
54        match self.providers.complete_cheapest(prompt, None, None).await {
55            Ok(response) => Ok(response),
56            Err(e) => {
57                tracing::warn!("LLM call failed for layer generation: {}", e);
58                Err(MemoryError::Embedding(format!(
59                    "failed to generate context layer: {}",
60                    e
61                )))
62            }
63        }
64    }
65}
66
67pub fn layer_type_token_limit(layer_type: LayerType) -> usize {
68    match layer_type {
69        LayerType::L0 => 150,
70        LayerType::L1 => 2500,
71        LayerType::L2 => 8000,
72    }
73}
74
75pub fn truncate_to_layer(content: &str, layer_type: LayerType) -> String {
76    let limit = layer_type_token_limit(layer_type);
77    let char_limit = limit * 4;
78
79    if content.len() <= char_limit {
80        return content.to_string();
81    }
82
83    let truncated = &content[..char_limit];
84    if let Some(last_newline) = truncated.rfind('\n') {
85        format!("{}...", truncated[..last_newline].trim())
86    } else if let Some(last_space) = truncated.rfind(' ') {
87        format!("{}...", truncated[..last_space].trim())
88    } else {
89        format!("{}...", truncated.trim())
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn test_layer_token_limits() {
99        assert_eq!(layer_type_token_limit(LayerType::L0), 150);
100        assert_eq!(layer_type_token_limit(LayerType::L1), 2500);
101        assert_eq!(layer_type_token_limit(LayerType::L2), 8000);
102    }
103
104    #[test]
105    fn test_truncate_to_layer() {
106        let long_text = "a".repeat(10000);
107        let truncated = truncate_to_layer(&long_text, LayerType::L0);
108        assert!(truncated.len() < 10000);
109        assert!(truncated.ends_with("..."));
110    }
111
112    #[test]
113    fn test_truncate_preserves_short() {
114        let short_text = "This is a short text.";
115        let result = truncate_to_layer(short_text, LayerType::L1);
116        assert_eq!(result, short_text);
117    }
118}