1use std::path::Path;
2use std::sync::Arc;
3
4use parking_lot::Mutex;
5use zkr::{
6 ClaimId, ClaimInput, ClaimKind, CorrectInput, DeleteInput, EmbeddingTarget, GetInput, MemoryDb,
7 PersonId, ProfilesInput, RememberInput, Remembered, RetrievalItem, RetrievalPack, SearchInput,
8 SourceId, SourceKind, TenantId,
9};
10
11pub struct ZkrStore {
12 db: Arc<Mutex<MemoryDb>>,
13 tenant_id: TenantId,
14 person_id: PersonId,
15}
16
17impl ZkrStore {
18 pub fn open(
19 path: &Path,
20 tenant_id: impl Into<String>,
21 person_id: impl Into<String>,
22 ) -> anyhow::Result<Self> {
23 if let Some(parent) = path.parent() {
24 std::fs::create_dir_all(parent)?;
25 }
26 Ok(Self {
27 db: Arc::new(Mutex::new(MemoryDb::open(path)?)),
28 tenant_id: TenantId::new(tenant_id.into())?,
29 person_id: PersonId::new(person_id.into())?,
30 })
31 }
32
33 pub async fn remember(
34 &self,
35 text: String,
36 kind: SourceKind,
37 ingestion_key: Option<String>,
38 claim: Option<ClaimInput>,
39 captured_at: i64,
40 ) -> anyhow::Result<Remembered> {
41 let db = Arc::clone(&self.db);
42 let tenant_id = self.tenant_id.clone();
43 let person_id = self.person_id.clone();
44 Ok(tokio::task::spawn_blocking(move || {
45 db.lock().remember(RememberInput {
46 tenant_id,
47 person_id,
48 ingestion_key,
49 kind,
50 text,
51 captured_at,
52 recorded_at: chrono::Utc::now().timestamp(),
53 claim,
54 feature_flag: None,
55 })
56 })
57 .await??)
58 }
59
60 pub async fn search(&self, query: String, limit: u32) -> anyhow::Result<RetrievalPack> {
61 let db = Arc::clone(&self.db);
62 let tenant_id = self.tenant_id.clone();
63 let person_id = self.person_id.clone();
64 Ok(tokio::task::spawn_blocking(move || {
65 db.lock().search(SearchInput {
66 tenant_id,
67 person_id,
68 query,
69 limit,
70 query_embedding: None,
71 as_of: None,
72 enabled_features: Vec::new(),
73 })
74 })
75 .await??)
76 }
77
78 pub async fn get(&self, kind: &str, id: String) -> anyhow::Result<RetrievalItem> {
79 let target = match kind {
80 "source" => EmbeddingTarget::Source(SourceId::new(id)?),
81 "evidence" => EmbeddingTarget::Evidence(zkr::EvidenceId::new(id)?),
82 "claim" => EmbeddingTarget::Claim(ClaimId::new(id)?),
83 _ => anyhow::bail!("kind must be source, evidence, or claim"),
84 };
85 let db = Arc::clone(&self.db);
86 let tenant_id = self.tenant_id.clone();
87 let person_id = self.person_id.clone();
88 Ok(tokio::task::spawn_blocking(move || {
89 db.lock().get(GetInput {
90 tenant_id,
91 person_id,
92 target,
93 })
94 })
95 .await??)
96 }
97
98 pub async fn correct(
99 &self,
100 claim_id: String,
101 text: String,
102 value: String,
103 valid_at: i64,
104 ) -> anyhow::Result<zkr::Corrected> {
105 let claim_id = ClaimId::new(claim_id).map_err(|e| anyhow::anyhow!(e))?;
106 let db = Arc::clone(&self.db);
107 let tenant_id = self.tenant_id.clone();
108 let person_id = self.person_id.clone();
109 Ok(tokio::task::spawn_blocking(move || {
110 db.lock().correct(CorrectInput {
111 tenant_id,
112 person_id,
113 claim_id,
114 text,
115 value,
116 valid_at,
117 recorded_at: chrono::Utc::now().timestamp(),
118 })
119 })
120 .await??)
121 }
122
123 pub async fn delete(&self, source_id: String) -> anyhow::Result<zkr::Deleted> {
124 let source_id = SourceId::new(source_id).map_err(|e| anyhow::anyhow!(e))?;
125 let db = Arc::clone(&self.db);
126 let tenant_id = self.tenant_id.clone();
127 let person_id = self.person_id.clone();
128 Ok(tokio::task::spawn_blocking(move || {
129 db.lock().delete_source(DeleteInput {
130 tenant_id,
131 person_id,
132 source_id,
133 deleted_at: chrono::Utc::now().timestamp(),
134 })
135 })
136 .await??)
137 }
138
139 pub async fn profiles(&self, limit: u32) -> anyhow::Result<Vec<zkr::ProfileEntry>> {
140 let db = Arc::clone(&self.db);
141 let tenant_id = self.tenant_id.clone();
142 let person_id = self.person_id.clone();
143 Ok(tokio::task::spawn_blocking(move || {
144 db.lock().profiles(ProfilesInput {
145 tenant_id,
146 person_id,
147 limit,
148 })
149 })
150 .await??)
151 }
152
153 pub async fn capture_turn(
154 &self,
155 chat_id: &str,
156 message_id: &str,
157 user: &str,
158 assistant: &str,
159 captured_at: i64,
160 ) -> anyhow::Result<()> {
161 self.remember(
162 user.to_string(),
163 SourceKind::Conversation,
164 Some(format!("{chat_id}:{message_id}:user")),
165 None,
166 captured_at,
167 )
168 .await?;
169 self.remember(
170 assistant.to_string(),
171 SourceKind::Conversation,
172 Some(format!("{chat_id}:{message_id}:assistant")),
173 None,
174 captured_at,
175 )
176 .await?;
177 Ok(())
178 }
179
180 pub async fn context(&self, query: &str, limit: u32) -> anyhow::Result<Option<String>> {
181 let pack = self.search(query.to_string(), limit).await?;
182 if pack.items.is_empty() {
183 return Ok(None);
184 }
185 let mut output = String::from("<zkr_memory>\n");
186 for item in pack.items {
187 let memory = serde_json::to_string(&item.memory)?;
188 let evidence = item
189 .evidence_ids
190 .iter()
191 .map(ToString::to_string)
192 .collect::<Vec<_>>()
193 .join(",");
194 output.push_str(&format!(
195 "- [{memory}] {} (evidence: {evidence}, relevance: {})\n",
196 item.excerpt, item.relevance_basis_points
197 ));
198 }
199 output.push_str("</zkr_memory>");
200 Ok(Some(output))
201 }
202
203 pub async fn record_reflection(
204 &self,
205 context: &str,
206 action: &str,
207 outcome: &str,
208 lesson: &str,
209 ) -> anyhow::Result<Remembered> {
210 let text =
211 format!("Context: {context}\nAction: {action}\nOutcome: {outcome}\nLesson: {lesson}");
212 let claim = ClaimInput {
213 subject: context.to_string(),
214 predicate: "improved".to_string(),
215 value: lesson.to_string(),
216 kind: ClaimKind::Skill,
217 valid_from: chrono::Utc::now().timestamp(),
218 tier: zkr::MemoryTier::LongTerm,
219 processing_state: zkr::MemoryProcessingState::Processed,
220 };
221 self.remember(
222 text,
223 SourceKind::Integration,
224 Some(format!("self_improve:{}", nanos())),
225 Some(claim),
226 chrono::Utc::now().timestamp(),
227 )
228 .await
229 }
230
231 pub async fn augment_prompt(&self, query: &str, base: &str) -> anyhow::Result<String> {
232 let mut lessons = Vec::new();
233 let mut seen = std::collections::HashSet::new();
234
235 for q in ["self improve lessons", query] {
236 let pack = self.search(q.to_string(), 5).await?;
237 for item in pack.items {
238 let text = item.excerpt.trim().to_string();
239 if seen.insert(text.clone()) {
240 lessons.push(text);
241 }
242 }
243 }
244
245 if lessons.is_empty() {
246 return Ok(base.to_string());
247 }
248
249 let mut augmented = base.to_string();
250 augmented.push_str("\n\n<lessons_learned>\n");
251 for (i, lesson) in lessons.iter().take(5).enumerate() {
252 augmented.push_str(&format!("{}. {lesson}\n", i + 1));
253 }
254 augmented.push_str("</lessons_learned>");
255 Ok(augmented)
256 }
257}
258
259fn nanos() -> u128 {
260 std::time::SystemTime::now()
261 .duration_since(std::time::UNIX_EPOCH)
262 .unwrap_or_default()
263 .as_nanos()
264}
265
266pub fn source_kind(value: &str) -> anyhow::Result<SourceKind> {
267 match value {
268 "conversation" => Ok(SourceKind::Conversation),
269 "screen" => Ok(SourceKind::Screen),
270 "audio" => Ok(SourceKind::Audio),
271 "document" => Ok(SourceKind::Document),
272 "integration" => Ok(SourceKind::Integration),
273 "user_correction" => Ok(SourceKind::UserCorrection),
274 _ => anyhow::bail!("invalid source kind"),
275 }
276}
277
278pub fn claim_kind(value: &str) -> anyhow::Result<ClaimKind> {
279 match value {
280 "fact" => Ok(ClaimKind::Fact),
281 "profile_fact" => Ok(ClaimKind::ProfileFact),
282 "preference" => Ok(ClaimKind::Preference),
283 "task" => Ok(ClaimKind::Task),
284 "skill" => Ok(ClaimKind::Skill),
285 "recommendation" => Ok(ClaimKind::Recommendation),
286 _ => anyhow::bail!("invalid claim kind"),
287 }
288}
289
290#[cfg(test)]
291mod tests {
292 use super::*;
293
294 #[tokio::test]
295 async fn remembers_and_retrieves_cited_source() {
296 let dir = tempfile::tempdir().unwrap();
297 let store = ZkrStore::open(&dir.path().join("memory.db"), "test", "person").unwrap();
298 let remembered = store
299 .remember(
300 "I prefer concise plans".to_string(),
301 SourceKind::Conversation,
302 Some("turn-1".to_string()),
303 None,
304 100,
305 )
306 .await
307 .unwrap();
308 let pack = store.search("concise plans".to_string(), 5).await.unwrap();
309 assert_eq!(pack.items.len(), 1);
310 assert_eq!(pack.items[0].evidence_ids, vec![remembered.evidence_id]);
311 }
312
313 #[tokio::test]
314 async fn ingestion_key_is_idempotent() {
315 let dir = tempfile::tempdir().unwrap();
316 let store = ZkrStore::open(&dir.path().join("memory.db"), "test", "person").unwrap();
317 let first = store
318 .remember(
319 "same turn".to_string(),
320 SourceKind::Conversation,
321 Some("turn-1".to_string()),
322 None,
323 100,
324 )
325 .await
326 .unwrap();
327 let second = store
328 .remember(
329 "same turn".to_string(),
330 SourceKind::Conversation,
331 Some("turn-1".to_string()),
332 None,
333 100,
334 )
335 .await
336 .unwrap();
337 assert_eq!(first.source_id, second.source_id);
338 }
339}