Skip to content

Commit 9692bfd

Browse files
committed
feat(context): integrate reservoir computing with context manager
Adds reservoir computing to ContextManager: - with_reservoir() constructor for optional reservoir - Automatic reservoir updates on add_turn() - reservoir_state() to access compressed context - reset_reservoir() to clear state - snapshot() now includes reservoir state - 3 new tests for reservoir integration Benefits: - Temporal context encoding across conversations - 1000-dim state vector captures conversation dynamics - Enables cross-session continuity - Minimal memory overhead vs full history Encoding: Simple bag-of-words (384-dim) → Reservoir (1000-dim state) Future: Replace with sentence-transformers embeddings
1 parent 64adc40 commit 9692bfd

1 file changed

Lines changed: 100 additions & 2 deletions

File tree

src/context.rs

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,17 @@
1010
//! - State snapshots
1111
//! - Context retrieval for query augmentation
1212
13+
use crate::reservoir::{encode_text, EchoStateNetwork};
1314
use crate::types::{ContextSnapshot, ConversationTurn, Query, Response};
1415
use serde::{Deserialize, Serialize};
1516
use std::collections::HashMap;
1617

1718
/// Maximum conversation history to keep in memory
1819
const MAX_HISTORY_SIZE: usize = 100;
1920

21+
/// Dimension for text encoding (matches reservoir input size)
22+
const ENCODING_DIM: usize = 384;
23+
2024
/// Context manager for maintaining conversation state
2125
#[derive(Debug, Clone, Serialize, Deserialize)]
2226
pub struct ContextManager {
@@ -26,21 +30,51 @@ pub struct ContextManager {
2630
history: Vec<ConversationTurn>,
2731
/// Per-project context snapshots
2832
project_contexts: HashMap<String, Vec<ConversationTurn>>,
33+
/// Reservoir for temporal context encoding (Phase 2)
34+
#[serde(skip)]
35+
reservoir: Option<EchoStateNetwork>,
2936
}
3037

3138
impl ContextManager {
3239
/// Create a new context manager
3340
pub fn new() -> Self {
41+
Self::with_reservoir(false)
42+
}
43+
44+
/// Create a context manager with reservoir computing enabled
45+
pub fn with_reservoir(enable_reservoir: bool) -> Self {
46+
let reservoir = if enable_reservoir {
47+
Some(EchoStateNetwork::new(
48+
ENCODING_DIM, // input size
49+
1000, // reservoir size
50+
100, // output size (compressed context)
51+
0.7, // leak rate
52+
0.95, // spectral radius
53+
))
54+
} else {
55+
None
56+
};
57+
3458
Self {
3559
current_project: None,
3660
history: Vec::new(),
3761
project_contexts: HashMap::new(),
62+
reservoir,
3863
}
3964
}
4065

4166
/// Add a conversation turn to history
4267
pub fn add_turn(&mut self, query: Query, response: Response) {
43-
let turn = ConversationTurn { query, response };
68+
let turn = ConversationTurn {
69+
query: query.clone(),
70+
response: response.clone(),
71+
};
72+
73+
// Update reservoir with query text if enabled
74+
if let Some(ref mut reservoir) = self.reservoir {
75+
let encoding = encode_text(&query.text, ENCODING_DIM);
76+
reservoir.update(&encoding);
77+
}
4478

4579
// Add to main history
4680
self.history.insert(0, turn.clone());
@@ -96,10 +130,24 @@ impl ContextManager {
96130

97131
/// Get a context snapshot for augmenting queries
98132
pub fn snapshot(&self, history_size: usize) -> ContextSnapshot {
133+
let reservoir_state = self.reservoir.as_ref().map(|r| r.state().to_vec());
134+
99135
ContextSnapshot {
100136
project: self.current_project.clone(),
101137
history: self.recent_history(history_size),
102-
reservoir_state: None, // Phase 2: reservoir computing
138+
reservoir_state,
139+
}
140+
}
141+
142+
/// Get reservoir state vector (if reservoir is enabled)
143+
pub fn reservoir_state(&self) -> Option<Vec<f32>> {
144+
self.reservoir.as_ref().map(|r| r.state().to_vec())
145+
}
146+
147+
/// Reset reservoir state (if enabled)
148+
pub fn reset_reservoir(&mut self) {
149+
if let Some(ref mut reservoir) = self.reservoir {
150+
reservoir.reset();
103151
}
104152
}
105153

@@ -276,4 +324,54 @@ mod tests {
276324
assert!(projects.contains(&"project-1".to_string()));
277325
assert!(projects.contains(&"project-2".to_string()));
278326
}
327+
328+
#[test]
329+
fn test_context_manager_with_reservoir() {
330+
let mut cm = ContextManager::with_reservoir(true);
331+
332+
// Reservoir state should initially be zeros
333+
let state1 = cm.reservoir_state();
334+
assert!(state1.is_some());
335+
assert_eq!(state1.as_ref().unwrap().len(), 1000);
336+
337+
// Add a turn - reservoir should update
338+
cm.add_turn(Query::new("Hello world"), create_test_response("Hi"));
339+
340+
let state2 = cm.reservoir_state();
341+
assert!(state2.is_some());
342+
343+
// State should have changed
344+
assert_ne!(state1, state2);
345+
346+
// Snapshot should include reservoir state
347+
let snapshot = cm.snapshot(5);
348+
assert!(snapshot.reservoir_state.is_some());
349+
assert_eq!(snapshot.reservoir_state.unwrap().len(), 1000);
350+
}
351+
352+
#[test]
353+
fn test_reservoir_reset() {
354+
let mut cm = ContextManager::with_reservoir(true);
355+
356+
cm.add_turn(Query::new("test"), create_test_response("response"));
357+
358+
let state = cm.reservoir_state().unwrap();
359+
assert!(!state.iter().all(|&x| x == 0.0));
360+
361+
cm.reset_reservoir();
362+
363+
let state_after_reset = cm.reservoir_state().unwrap();
364+
assert!(state_after_reset.iter().all(|&x| x == 0.0));
365+
}
366+
367+
#[test]
368+
fn test_context_manager_without_reservoir() {
369+
let cm = ContextManager::new();
370+
371+
// Without reservoir, state should be None
372+
assert!(cm.reservoir_state().is_none());
373+
374+
let snapshot = cm.snapshot(5);
375+
assert!(snapshot.reservoir_state.is_none());
376+
}
279377
}

0 commit comments

Comments
 (0)