tandem_memory/
context_layers.rs1use 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}