1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde::Deserialize;
5use zkr::{ClaimInput, MemoryProcessingState, MemoryTier};
6
7use super::{Tool, ToolResult, ToolSpec};
8use crate::memory::zkr::{claim_kind, source_kind, ZkrStore};
9
10pub struct ZkrTool {
11 store: Arc<ZkrStore>,
12}
13
14impl ZkrTool {
15 pub fn new(store: Arc<ZkrStore>) -> Self {
16 Self { store }
17 }
18}
19
20#[derive(Deserialize)]
21struct ZkrArgs {
22 action: String,
23 text: Option<String>,
24 source_kind: Option<String>,
25 ingestion_key: Option<String>,
26 subject: Option<String>,
27 predicate: Option<String>,
28 value: Option<String>,
29 claim_kind: Option<String>,
30 valid_from: Option<i64>,
31 query: Option<String>,
32 limit: Option<u32>,
33 kind: Option<String>,
34 id: Option<String>,
35 claim_id: Option<String>,
36 source_id: Option<String>,
37 valid_at: Option<i64>,
38}
39
40#[async_trait]
41impl Tool for ZkrTool {
42 fn name(&self) -> &str {
43 "zkr_memory"
44 }
45
46 fn spec(&self) -> ToolSpec {
47 ToolSpec {
48 name: self.name().to_string(),
49 description: "Store, search, inspect, correct, and delete evidence-backed temporal memory with citations.".to_string(),
50 parameters: serde_json::json!({
51 "type": "object",
52 "properties": {
53 "action": { "type": "string", "enum": ["remember", "search", "get", "correct", "delete", "profiles"] },
54 "text": { "type": "string" },
55 "source_kind": { "type": "string", "enum": ["conversation", "screen", "audio", "document", "integration", "user_correction"] },
56 "ingestion_key": { "type": "string" },
57 "subject": { "type": "string" },
58 "predicate": { "type": "string" },
59 "value": { "type": "string" },
60 "claim_kind": { "type": "string", "enum": ["fact", "profile_fact", "preference", "task", "skill", "recommendation"] },
61 "valid_from": { "type": "integer" },
62 "query": { "type": "string" },
63 "limit": { "type": "integer", "minimum": 1, "maximum": 100 },
64 "kind": { "type": "string", "enum": ["source", "evidence", "claim"] },
65 "id": { "type": "string" },
66 "claim_id": { "type": "string" },
67 "source_id": { "type": "string" },
68 "valid_at": { "type": "integer" }
69 },
70 "required": ["action"]
71 }),
72 }
73 }
74
75 async fn execute(&self, arguments: &str) -> anyhow::Result<ToolResult> {
76 let args: ZkrArgs = serde_json::from_str(arguments)?;
77 let result = match args.action.as_str() {
78 "remember" => {
79 let text = args
80 .text
81 .ok_or_else(|| anyhow::anyhow!("text is required"))?;
82 let claim = match (args.subject, args.predicate, args.value, args.claim_kind) {
83 (Some(subject), Some(predicate), Some(value), Some(kind)) => Some(ClaimInput {
84 subject,
85 predicate,
86 value,
87 kind: claim_kind(&kind)?,
88 valid_from: args
89 .valid_from
90 .unwrap_or_else(|| chrono::Utc::now().timestamp()),
91 tier: MemoryTier::LongTerm,
92 processing_state: MemoryProcessingState::Processed,
93 }),
94 (None, None, None, None) => None,
95 _ => anyhow::bail!(
96 "subject, predicate, value, and claim_kind must be provided together"
97 ),
98 };
99 let remembered = self
100 .store
101 .remember(
102 text,
103 source_kind(args.source_kind.as_deref().unwrap_or("integration"))?,
104 args.ingestion_key,
105 claim,
106 chrono::Utc::now().timestamp(),
107 )
108 .await?;
109 serde_json::to_string_pretty(&remembered)?
110 }
111 "search" => serde_json::to_string_pretty(
112 &self
113 .store
114 .search(
115 args.query
116 .ok_or_else(|| anyhow::anyhow!("query is required"))?,
117 args.limit.unwrap_or(10),
118 )
119 .await?,
120 )?,
121 "get" => serde_json::to_string_pretty(
122 &self
123 .store
124 .get(
125 args.kind
126 .as_deref()
127 .ok_or_else(|| anyhow::anyhow!("kind is required"))?,
128 args.id.ok_or_else(|| anyhow::anyhow!("id is required"))?,
129 )
130 .await?,
131 )?,
132 "correct" => serde_json::to_string_pretty(
133 &self
134 .store
135 .correct(
136 args.claim_id
137 .ok_or_else(|| anyhow::anyhow!("claim_id is required"))?,
138 args.text
139 .ok_or_else(|| anyhow::anyhow!("text is required"))?,
140 args.value
141 .ok_or_else(|| anyhow::anyhow!("value is required"))?,
142 args.valid_at
143 .ok_or_else(|| anyhow::anyhow!("valid_at is required"))?,
144 )
145 .await?,
146 )?,
147 "delete" => serde_json::to_string_pretty(
148 &self
149 .store
150 .delete(
151 args.source_id
152 .ok_or_else(|| anyhow::anyhow!("source_id is required"))?,
153 )
154 .await?,
155 )?,
156 "profiles" => {
157 serde_json::to_string_pretty(&self.store.profiles(args.limit.unwrap_or(20)).await?)?
158 }
159 other => anyhow::bail!("unknown zkr action: {other}"),
160 };
161 Ok(ToolResult::success(result))
162 }
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168
169 #[tokio::test]
170 async fn remember_and_search_actions_return_citations() {
171 let dir = tempfile::tempdir().unwrap();
172 let store =
173 Arc::new(ZkrStore::open(&dir.path().join("memory.db"), "test", "person").unwrap());
174 let tool = ZkrTool::new(store);
175 let remembered = tool
176 .execute(r#"{"action":"remember","text":"prefers short answers","source_kind":"conversation"}"#)
177 .await
178 .unwrap();
179 assert!(!remembered.is_error);
180 let searched = tool
181 .execute(r#"{"action":"search","query":"short answers"}"#)
182 .await
183 .unwrap();
184 assert!(searched.output.contains("evidence_ids"));
185 }
186}