cascade_agent/memory/
mod.rs

1//! Memory & Context Management.
2//!
3//! Provides token counting, context compaction (summarization), and orchestration
4//! via [`MemoryManager`].
5
6pub mod compaction;
7pub mod state;
8pub mod tokenizer;
9pub mod types;
10
11use std::sync::Arc;
12
13use crate::config::MemorySettings;
14use crate::error::Result;
15use crate::memory::compaction::Compactor;
16use crate::memory::state::ConversationStateAccess;
17use crate::memory::tokenizer::TokenCounter;
18use crate::memory::types::CompactionReport;
19
20/// Orchestrates token counting and context compaction for conversation management.
21///
22/// `MemoryManager` is the main entry point for memory operations. It checks whether
23/// compaction is needed and delegates to the [`Compactor`] when the conversation
24/// exceeds configured token limits.
25pub struct MemoryManager {
26    token_counter: TokenCounter,
27    compactor: Compactor,
28    context_token_limit: usize,
29}
30
31impl MemoryManager {
32    /// Create a new memory manager.
33    ///
34    /// # Arguments
35    /// * `settings` - Memory configuration (token limits, compaction ratio, etc.)
36    /// * `cascade_config` - llm-cascade `AppConfig` for summarization inference
37    /// * `db_conn` - SQLite connection for llm-cascade cooldown/logging
38    pub fn new(
39        settings: &MemorySettings,
40        cascade_config: Arc<llm_cascade::AppConfig>,
41        db_conn: rusqlite::Connection,
42    ) -> Result<Self> {
43        let token_counter = TokenCounter::new(&settings.tokenizer_model)?;
44        let compactor = Compactor::new(settings, cascade_config, db_conn);
45
46        Ok(Self {
47            token_counter,
48            compactor,
49            context_token_limit: settings.context_token_limit,
50        })
51    }
52
53    /// Count the total tokens in a conversation state.
54    pub fn count_tokens(&self, state: &impl ConversationStateAccess) -> usize {
55        self.token_counter.count_messages(state.messages())
56    }
57
58    /// Check whether compaction should be triggered based on current token count.
59    ///
60    /// Compaction is recommended when tokens exceed the configured limit.
61    pub fn should_compact(&self, token_count: usize) -> bool {
62        token_count > self.context_token_limit
63    }
64
65    /// Attempt to compact the conversation if it exceeds the token limit.
66    ///
67    /// If compaction is not needed, returns a no-op report.
68    /// If compaction fails (e.g., summarization LLM error), returns the error
69    /// without modifying the conversation state.
70    pub async fn compact(
71        &self,
72        state: &mut impl ConversationStateAccess,
73    ) -> Result<CompactionReport> {
74        let tokens = self.count_tokens(state);
75
76        if !self.should_compact(tokens) {
77            return Ok(CompactionReport {
78                messages_before: state.messages().len(),
79                messages_after: state.messages().len(),
80                tokens_before: tokens,
81                tokens_after: tokens,
82            });
83        }
84
85        self.compactor
86            .compact(state, &self.token_counter, self.context_token_limit)
87            .await
88    }
89
90    /// Returns the configured context token limit.
91    pub fn token_limit(&self) -> usize {
92        self.context_token_limit
93    }
94
95    /// Returns a reference to the underlying token counter.
96    pub fn token_counter(&self) -> &TokenCounter {
97        &self.token_counter
98    }
99}
cascade_agent/memory/mod.rs

cascade_agent/memory/
mod.rs
