1pub use crate::utils::cosine_similarity;
12use anyhow::Result;
13use compact_str::CompactString;
14use serde_json::Value;
15use std::{future::Future, sync::Arc};
16
17mod embedder;
18mod mem;
19mod sqlite;
20pub mod tools;
21mod utils;
22
23pub use embedder::{Embedder, NoEmbedder};
24pub use mem::InMemory;
25pub use sqlite::SqliteMemory;
26
27#[derive(Debug, Clone, Default)]
29pub struct MemoryEntry {
30 pub key: CompactString,
32 pub value: String,
34 pub metadata: Option<Value>,
36 pub created_at: u64,
38 pub accessed_at: u64,
40 pub access_count: u32,
42 pub embedding: Option<Vec<f32>>,
44}
45
46#[derive(Debug, Clone, Default)]
48pub struct RecallOptions {
49 pub limit: usize,
51 pub time_range: Option<(u64, u64)>,
53 pub relevance_threshold: Option<f32>,
55}
56
57pub trait Memory: Send + Sync {
65 fn get(&self, key: &str) -> Option<String>;
67
68 fn entries(&self) -> Vec<(String, String)>;
70
71 fn set(&self, key: impl Into<String>, value: impl Into<String>) -> Option<String>;
73
74 fn remove(&self, key: &str) -> Option<String>;
76
77 fn compile(&self) -> String {
79 let entries = self.entries();
80 if entries.is_empty() {
81 return String::new();
82 }
83
84 let mut out = String::from("<memory>\n");
85 for (key, value) in &entries {
86 out.push_str(&format!("<{key}>\n"));
87 out.push_str(value);
88 if !value.ends_with('\n') {
89 out.push('\n');
90 }
91 out.push_str(&format!("</{key}>\n"));
92 }
93 out.push_str("</memory>");
94 out
95 }
96
97 fn store(
99 &self,
100 key: impl Into<String> + Send,
101 value: impl Into<String> + Send,
102 ) -> impl Future<Output = Result<()>> + Send {
103 self.set(key, value);
104 async { Ok(()) }
105 }
106
107 fn recall(
109 &self,
110 _query: &str,
111 _options: RecallOptions,
112 ) -> impl Future<Output = Result<Vec<MemoryEntry>>> + Send {
113 async { Ok(Vec::new()) }
114 }
115
116 fn compile_relevant(&self, _query: &str) -> impl Future<Output = String> + Send {
118 let compiled = self.compile();
119 async move { compiled }
120 }
121}
122
123pub fn with_memory(mut config: wcore::AgentConfig, memory: &impl Memory) -> wcore::AgentConfig {
125 let compiled = memory.compile();
126 if !compiled.is_empty() {
127 config.system_prompt = format!("{}\n\n{compiled}", config.system_prompt);
128 }
129 config
130}
131
132impl wcore::Hook for InMemory {
133 fn on_build_agent(&self, config: wcore::AgentConfig) -> wcore::AgentConfig {
134 with_memory(config, self)
135 }
136
137 fn on_register_tools(
138 &self,
139 registry: &mut wcore::ToolRegistry,
140 ) -> impl std::future::Future<Output = ()> + Send {
141 let mem = Arc::new(self.clone());
142 let remember = tools::remember(Arc::clone(&mem));
143 let recall = tools::recall(mem);
144 registry.insert(remember.tool, remember.handler);
145 registry.insert(recall.tool, recall.handler);
146 async {}
147 }
148}