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