1010//! - State snapshots
1111//! - Context retrieval for query augmentation
1212
13+ use crate :: reservoir:: { encode_text, EchoStateNetwork } ;
1314use crate :: types:: { ContextSnapshot , ConversationTurn , Query , Response } ;
1415use serde:: { Deserialize , Serialize } ;
1516use std:: collections:: HashMap ;
1617
1718/// Maximum conversation history to keep in memory
1819const 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 ) ]
2226pub 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
3138impl 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