1mod causal;
18mod meta;
19mod online;
20mod reflexion;
21mod sessions;
22mod skills;
23mod temporal;
24
25pub use causal::{CausalEdge, CausalMemory, CausalNode, Hyperedge};
26pub use meta::{FewShotLearner, LearningStrategy, MetaLearner, TaskFeatures};
27pub use online::{DriftDetector, Experience, ExperienceWindow, OnlineLearner, ParameterSnapshot};
28pub use reflexion::{Critique, CritiqueType, ReflexionEpisode, ReflexionMemory};
29pub use sessions::{LearningSession, Reward, SessionTurn};
30pub use skills::{Skill, SkillLibrary, SkillPattern};
31pub use temporal::{TemporalMemory, TemporalOccurrence, TemporalPeriod, TimeCrystal};
32
33use serde::{Deserialize, Serialize};
34use std::sync::Arc;
35use tokio::sync::RwLock;
36use uuid::Uuid;
37
38pub struct AgenticDB {
49 pub reflexion: Arc<RwLock<ReflexionMemory>>,
51
52 pub skills: Arc<RwLock<SkillLibrary>>,
54
55 pub causal: Arc<RwLock<CausalMemory>>,
57
58 pub sessions: Arc<RwLock<Vec<LearningSession>>>,
60
61 pub temporal: Arc<RwLock<TemporalMemory>>,
63
64 pub meta: Arc<RwLock<MetaLearner>>,
66
67 pub experiences: Arc<RwLock<ExperienceWindow>>,
69
70 pub drift_detector: Arc<RwLock<DriftDetector>>,
72
73 #[allow(dead_code)]
75 path: Option<String>,
76}
77
78impl AgenticDB {
79 pub fn new() -> Self {
85 Self::with_dimensions(64, 64)
86 }
87
88 pub fn with_dimensions(meta_param_dim: usize, drift_feature_dim: usize) -> Self {
90 Self {
91 reflexion: Arc::new(RwLock::new(ReflexionMemory::new())),
92 skills: Arc::new(RwLock::new(SkillLibrary::new())),
93 causal: Arc::new(RwLock::new(CausalMemory::new())),
94 sessions: Arc::new(RwLock::new(Vec::new())),
95 temporal: Arc::new(RwLock::new(TemporalMemory::new())),
96 meta: Arc::new(RwLock::new(MetaLearner::new("agent_meta", meta_param_dim))),
97 experiences: Arc::new(RwLock::new(ExperienceWindow::new(1000))),
98 drift_detector: Arc::new(RwLock::new(DriftDetector::new(drift_feature_dim))),
99 path: None,
100 }
101 }
102
103 pub fn with_path(path: impl Into<String>) -> Self {
105 Self {
106 reflexion: Arc::new(RwLock::new(ReflexionMemory::new())),
107 skills: Arc::new(RwLock::new(SkillLibrary::new())),
108 causal: Arc::new(RwLock::new(CausalMemory::new())),
109 sessions: Arc::new(RwLock::new(Vec::new())),
110 temporal: Arc::new(RwLock::new(TemporalMemory::new())),
111 meta: Arc::new(RwLock::new(MetaLearner::new("agent_meta", 64))),
112 experiences: Arc::new(RwLock::new(ExperienceWindow::new(1000))),
113 drift_detector: Arc::new(RwLock::new(DriftDetector::new(64))),
114 path: Some(path.into()),
115 }
116 }
117
118 pub async fn add_reflexion(&self, episode: ReflexionEpisode) {
144 let mut reflexion = self.reflexion.write().await;
145 reflexion.add_episode(episode);
146 }
147
148 pub async fn query_similar_failures(&self, task: &str, limit: usize) -> Vec<ReflexionEpisode> {
150 let reflexion = self.reflexion.read().await;
151 reflexion.find_similar_failures(task, limit)
152 }
153
154 pub async fn register_skill(&self, skill: Skill) {
156 let mut skills = self.skills.write().await;
157 skills.add_skill(skill);
158 }
159
160 pub async fn find_skills(&self, _task_description: &str) -> Vec<&Skill> {
162 Vec::new()
164 }
165
166 pub async fn add_causal_link(
168 &self,
169 cause: Uuid,
170 effect: Uuid,
171 relationship: impl Into<String>,
172 strength: f32,
173 ) {
174 let mut causal = self.causal.write().await;
175 causal.add_edge(CausalEdge {
176 id: Uuid::new_v4(),
177 cause,
178 effect,
179 relationship: relationship.into(),
180 strength,
181 evidence_count: 1,
182 });
183 }
184
185 pub async fn start_session(&self, goal: impl Into<String>) -> Uuid {
187 let session = LearningSession::new(goal);
188 let id = session.id;
189 let mut sessions = self.sessions.write().await;
190 sessions.push(session);
191 id
192 }
193
194 pub async fn record_turn(
196 &self,
197 session_id: Uuid,
198 action: impl Into<String>,
199 observation: impl Into<String>,
200 reward: Reward,
201 ) {
202 let mut sessions = self.sessions.write().await;
203 if let Some(session) = sessions.iter_mut().find(|s| s.id == session_id) {
204 session.add_turn(SessionTurn {
205 action: action.into(),
206 observation: observation.into(),
207 reward,
208 timestamp: chrono::Utc::now(),
209 });
210 }
211 }
212
213 pub async fn stats(&self) -> AgenticStats {
215 let reflexion = self.reflexion.read().await;
216 let skills = self.skills.read().await;
217 let causal = self.causal.read().await;
218 let sessions = self.sessions.read().await;
219 let temporal = self.temporal.read().await;
220 let meta = self.meta.read().await;
221 let experiences = self.experiences.read().await;
222
223 AgenticStats {
224 reflexion_episodes: reflexion.len(),
225 failed_episodes: reflexion.failure_count(),
226 skills_count: skills.len(),
227 causal_nodes: causal.node_count(),
228 causal_edges: causal.edge_count(),
229 total_sessions: sessions.len(),
230 total_turns: sessions.iter().map(|s| s.turns.len()).sum(),
231 temporal_patterns: temporal.len(),
232 meta_tasks_learned: meta.num_tasks() as usize,
233 learning_strategies: meta.num_strategies(),
234 experience_buffer_size: experiences.len(),
235 }
236 }
237
238 pub async fn record_temporal(&self, pattern_name: &str, period: TemporalPeriod, value: f32) {
245 let mut temporal = self.temporal.write().await;
246 temporal.record(pattern_name, period, value);
247 }
248
249 pub async fn predict_temporal(&self, pattern_name: &str) -> Option<f32> {
251 let temporal = self.temporal.read().await;
252 temporal.predict(pattern_name)
253 }
254
255 pub async fn add_experience(
257 &mut self,
258 features: Vec<f32>,
259 target: f32,
260 task_id: Option<String>,
261 ) {
262 let mut experiences = self.experiences.write().await;
263 experiences.add(features, target, task_id);
264 }
265
266 pub async fn check_drift(&self, features: &[f32]) -> bool {
271 let mut detector = self.drift_detector.write().await;
272 detector.update(features)
273 }
274
275 pub async fn meta_update(
277 &self,
278 task_id: &str,
279 final_params: &[f32],
280 task_embedding: Option<Vec<f32>>,
281 ) {
282 let mut meta = self.meta.write().await;
283 meta.meta_update(task_id, final_params, task_embedding);
284 }
285
286 pub async fn get_task_initialization(&self, task_embedding: Option<&[f32]>) -> Vec<f32> {
288 let meta = self.meta.read().await;
289 meta.initialize_for_task(task_embedding)
290 }
291
292 pub async fn register_strategy(&self, strategy: LearningStrategy) {
294 let mut meta = self.meta.write().await;
295 meta.register_strategy(strategy);
296 }
297
298 pub async fn select_strategy(&self, task_features: &TaskFeatures) -> Option<String> {
300 let meta = self.meta.read().await;
301 meta.select_strategy(task_features).map(|s| s.name.clone())
302 }
303}
304
305impl Default for AgenticDB {
306 fn default() -> Self {
307 Self::new()
308 }
309}
310
311#[derive(Debug, Clone, Serialize, Deserialize)]
313pub struct AgenticStats {
314 pub reflexion_episodes: usize,
316 pub failed_episodes: usize,
318 pub skills_count: usize,
320 pub causal_nodes: usize,
322 pub causal_edges: usize,
324 pub total_sessions: usize,
326 pub total_turns: usize,
328 pub temporal_patterns: usize,
330 pub meta_tasks_learned: usize,
332 pub learning_strategies: usize,
334 pub experience_buffer_size: usize,
336}
337
338#[cfg(test)]
339mod tests {
340 use super::*;
341
342 #[tokio::test]
351 async fn test_reflexion_workflow() {
352 let db = AgenticDB::new();
354
355 let episode = ReflexionEpisode::new(
357 "code_generation", "Write a function to calculate factorial", "fn factorial(n: i32) -> i32 { n * factorial(n-1) }", false, )
362 .with_critique(Critique::new(
363 CritiqueType::LogicError,
364 "Missing base case causes infinite recursion", "Add: if n <= 1 { return 1; }", ))
367 .with_critique(Critique::new(
368 CritiqueType::MissingStep,
369 "No handling for negative numbers",
370 "Add input validation or use unsigned type",
371 ));
372
373 db.add_reflexion(episode).await;
375
376 let stats = db.stats().await;
378 assert_eq!(stats.reflexion_episodes, 1);
379 assert_eq!(stats.failed_episodes, 1);
380
381 let similar = db.query_similar_failures("factorial function", 5).await;
383 assert_eq!(similar.len(), 1);
384
385 let past_mistake = &similar[0];
387 assert!(!past_mistake.succeeded);
388 assert_eq!(past_mistake.critiques.len(), 2);
389 }
390
391 #[tokio::test]
398 async fn test_skill_library() {
399 let db = AgenticDB::new();
400
401 let skill = Skill::new(
403 "error_handling",
404 "Rust Error Handling Pattern",
405 vec![
406 SkillPattern::new("result_type", "fn do_thing() -> Result<T, Error> { ... }"),
407 SkillPattern::new("question_mark", "let value = risky_op()?;"),
408 ],
409 )
410 .with_success_rate(0.95)
411 .with_usage_count(42);
412
413 db.register_skill(skill).await;
414
415 let stats = db.stats().await;
416 assert_eq!(stats.skills_count, 1);
417 }
418
419 #[tokio::test]
425 async fn test_causal_memory() {
426 let db = AgenticDB::new();
427
428 let cause_id = Uuid::new_v4();
429 let effect_id = Uuid::new_v4();
430
431 db.add_causal_link(
433 cause_id, effect_id, "causes", 0.8, )
435 .await;
436
437 let stats = db.stats().await;
438 assert_eq!(stats.causal_edges, 1);
439 }
440
441 #[tokio::test]
449 async fn test_learning_session() {
450 let db = AgenticDB::new();
451
452 let session_id = db.start_session("Fix the bug in auth module").await;
454
455 db.record_turn(
457 session_id,
458 "read_file auth.rs",
459 "Found potential null check issue on line 42",
460 Reward::Neutral,
461 )
462 .await;
463
464 db.record_turn(
466 session_id,
467 "edit auth.rs: add None check",
468 "File updated successfully",
469 Reward::Positive(0.5),
470 )
471 .await;
472
473 db.record_turn(
475 session_id,
476 "run tests",
477 "All 15 tests passing",
478 Reward::Positive(1.0), )
480 .await;
481
482 let stats = db.stats().await;
483 assert_eq!(stats.total_sessions, 1);
484 assert_eq!(stats.total_turns, 3);
485 }
486
487 #[tokio::test]
496 async fn test_integrated_workflow() {
497 let db = AgenticDB::new();
498
499 let past_failure = ReflexionEpisode::new(
501 "api_design",
502 "Design REST API endpoint",
503 "POST /users with no validation",
504 false,
505 )
506 .with_critique(Critique::new(
507 CritiqueType::DesignFlaw,
508 "No input validation leads to security issues",
509 "Always validate and sanitize inputs",
510 ));
511 db.add_reflexion(past_failure).await;
512
513 let skill = Skill::new(
515 "api_validation",
516 "Input Validation Pattern",
517 vec![SkillPattern::new("validate_first", "validate(&input)?;")],
518 );
519 db.register_skill(skill).await;
520
521 let validation_id = Uuid::new_v4();
523 let security_id = Uuid::new_v4();
524 db.add_causal_link(validation_id, security_id, "improves", 0.9)
525 .await;
526
527 let session_id = db
529 .start_session("Implement user registration endpoint")
530 .await;
531
532 let failures = db.query_similar_failures("api endpoint", 5).await;
534 assert!(!failures.is_empty(), "Should find past API failure");
535
536 db.record_turn(
538 session_id,
539 "apply validation skill",
540 "Added input validation with proper error handling",
541 Reward::Positive(0.8),
542 )
543 .await;
544
545 let stats = db.stats().await;
547 assert_eq!(stats.reflexion_episodes, 1);
548 assert_eq!(stats.skills_count, 1);
549 assert_eq!(stats.causal_edges, 1);
550 assert_eq!(stats.total_sessions, 1);
551 assert_eq!(stats.total_turns, 1);
552
553 println!("AgenticDB Stats: {:?}", stats);
554 }
555}