1use crate::{storage::{self, KnowledgeBase}, text, types::*, Error, Result};
2use rusqlite::{params, Connection};
3use serde::{Deserialize, Serialize};
4use serde_json::json;
5use std::collections::BTreeSet;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8#[serde(default)]
9pub struct MemoryState {
10 pub pinned: bool,
11 pub strength: i64,
12 pub useful_count: i64,
13 pub useful_score: f64,
14 pub last_recalled_at_us: Option<i64>,
15 pub last_decay_at_us: Option<i64>,
16}
17impl Default for MemoryState {
18 fn default() -> Self { Self { pinned: false, strength: 1, useful_count: 0, useful_score: 0.0, last_recalled_at_us: None, last_decay_at_us: None } }
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct Memory {
23 #[serde(flatten)] pub header: RecordHeader,
24 pub memory_type: String,
25 pub judgment: String,
26 pub reasoning: String,
27 pub state: MemoryState,
28}
29
30fn knowledge() -> String { "knowledge".into() }
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct MemoryInput {
33 #[serde(flatten)] pub record: RecordInput,
34 #[serde(default = "knowledge")] pub memory_type: String,
35 pub judgment: String,
36 #[serde(default)] pub reasoning: String,
37 #[serde(default)] pub state: Option<MemoryState>,
39}
40impl MemoryInput {
41 pub fn new(judgment: impl Into<String>) -> Self {
42 Self { record: RecordInput::default(), memory_type: knowledge(), judgment: judgment.into(), reasoning: String::new(), state: None }
43 }
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47#[serde(default)]
48pub struct DecayPolicy {
49 pub tier0_threshold: f64,
50 pub tier1_threshold: f64,
51 pub useful_score_boost: f64,
52 pub strength_boost: i64,
53 pub tier0_cycle_days: u32,
54}
55impl Default for DecayPolicy {
56 fn default() -> Self { Self { tier0_threshold: 3.0, tier1_threshold: 10.0, useful_score_boost: 2.5, strength_boost: 1, tier0_cycle_days: 3 } }
57}
58impl DecayPolicy {
59 fn validate(&self) -> Result<()> {
60 if !self.tier0_threshold.is_finite() || !self.tier1_threshold.is_finite() || !self.useful_score_boost.is_finite()
61 || self.tier0_threshold <= 0.0 || self.tier1_threshold <= self.tier0_threshold
62 || self.useful_score_boost <= 0.0 || self.strength_boost < 1 || self.tier0_cycle_days == 0 {
63 return Err(Error::Validation("invalid decay policy".into()));
64 }
65 Ok(())
66 }
67 pub fn tier(&self, score: f64) -> u8 {
68 if score >= self.tier1_threshold { 2 } else if score >= self.tier0_threshold { 1 } else { 0 }
69 }
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize, Default)]
73pub struct FeedbackRequest {
74 #[serde(default)] pub filter: ReadFilter,
75 #[serde(default)] pub recalled_ids: Vec<i64>,
76 #[serde(default)] pub useful_ids: Vec<i64>,
77 #[serde(default)] pub now_us: Option<i64>,
78 #[serde(default)] pub policy: DecayPolicy,
79}
80#[derive(Debug, Clone, Default, Serialize, Deserialize)]
81pub struct FeedbackReport { pub recalled: usize, pub boosted: usize, pub penalized: usize }
82#[derive(Debug, Clone, Default, Serialize, Deserialize)]
83pub struct DecayReport { pub decayed: usize, pub retirement_candidates: Vec<RecordKey> }
84
85#[derive(Clone)]
86pub struct MemoryStore(pub(crate) KnowledgeBase);
87
88pub(crate) fn upsert(conn: &Connection, input: &MemoryInput) -> Result<(Memory, crate::index::IndexDocument)> {
89 storage::validate_identity("memory_type", &input.memory_type)?;
90 if input.judgment.trim().is_empty() { return Err(Error::Validation("judgment is required".into())); }
91 let existing: Option<Memory> = if let Some(id) = input.record.id {
92 storage::record_value(conn, &RecordKey { id })?.map(serde_json::from_value).transpose()?
93 } else { None };
94 let state = input.state.clone().or_else(|| existing.map(|m| m.state)).unwrap_or_default();
95 if state.strength < 0 || state.useful_count < 0 || !state.useful_score.is_finite() || state.useful_score < 0.0 {
96 return Err(Error::Validation("memory lifecycle values must be finite and nonnegative".into()));
97 }
98 let judgment = input.judgment.trim().to_string();
99 let reasoning = input.reasoning.trim().to_string();
100 let memory_type_id = storage::term_id(conn, &input.memory_type)?;
101 let payload = json!({"memory_type_id": memory_type_id, "judgment": judgment, "reasoning": reasoning,
102 "state": state, "judgment_key": text::normalized_tag(&judgment)});
103 let (header, document) = storage::put_record(conn, RecordKind::Memory, &input.record, &payload, &judgment)?;
104 Ok((Memory { header, memory_type: input.memory_type.clone(), judgment, reasoning, state }, document))
105}
106
107pub(crate) fn as_input(memory: &Memory, at: i64) -> MemoryInput {
108 MemoryInput {
109 record: RecordInput { id: Some(memory.header.id), namespace: memory.header.namespace.clone(),
110 scope: memory.header.scope.clone(), tags: memory.header.tags.clone(), evidence: memory.header.evidence.clone(),
111 metadata: memory.header.metadata.clone(), created_at_us: Some(memory.header.created_at_us),
112 updated_at_us: Some(at.max(memory.header.updated_at_us)), expected_revision: Some(memory.header.revision) },
113 memory_type: memory.memory_type.clone(), judgment: memory.judgment.clone(), reasoning: memory.reasoning.clone(), state: Some(memory.state.clone()),
114 }
115}
116
117impl MemoryStore {
118 pub fn upsert(&self, input: MemoryInput) -> Result<WriteReceipt<Memory>> {
119 let receipt = self.0.mutate(|tx| upsert(tx, &input))?;
120 let WriteReceipt { value: (memory, document), revision } = receipt;
121 self.0.index_documents(&[document])?;
124 Ok(WriteReceipt { value: memory, revision })
125 }
126 pub fn upsert_many(&self, inputs: &[MemoryInput]) -> Result<WriteReceipt<Vec<Memory>>> {
127 let receipt = self.0.mutate(|tx| inputs.iter().map(|input| upsert(tx, input)).collect::<Result<Vec<_>>>())?;
128 let WriteReceipt { value, revision } = receipt;
129 let (memories, documents): (Vec<Memory>, Vec<crate::index::IndexDocument>) = value.into_iter().unzip();
130 self.0.index_documents(&documents)?;
131 Ok(WriteReceipt { value: memories, revision })
132 }
133 pub fn upsert_by_judgment(&self, mut input: MemoryInput) -> Result<WriteReceipt<Memory>> {
134 let receipt = self.0.mutate(|tx| {
135 let mut stmt = tx.prepare("SELECT id FROM records WHERE namespace_id=(SELECT id FROM strings WHERE text=?1)
136 AND scope_id=(SELECT id FROM strings WHERE text=?2) AND kind=?3 AND json_extract(payload_json,'$.judgment_key')=?4 ORDER BY id LIMIT 2")?;
137 let ids = stmt.query_map(params![text::normalized_tag(&input.record.namespace), text::normalized_tag(&input.record.scope),
138 RecordKind::Memory.code(), text::normalized_tag(&input.judgment)], |r| r.get::<_, i64>(0))?.collect::<std::result::Result<Vec<_>, _>>()?;
139 if ids.len() > 1 { return Err(Error::Conflict("multiple memories have this judgment; update by ID".into())); }
140 if let Some(id) = ids.first() {
141 if input.record.id.is_some_and(|given| given != *id) { return Err(Error::Conflict("judgment belongs to a different ID".into())); }
142 input.record.id = Some(*id);
143 let key = RecordKey { id: *id };
144 let existing: Memory = serde_json::from_value(storage::record_value(tx, &key)?.ok_or_else(|| Error::NotFound(id.to_string()))?)?;
145 let mut metadata = existing.header.metadata;
146 metadata.extend(input.record.metadata.clone());
147 input.record.metadata = metadata;
148 if input.record.evidence.is_empty() { input.record.evidence = existing.header.evidence; }
149 }
150 upsert(tx, &input)
151 })?;
152 let WriteReceipt { value: (memory, document), revision } = receipt;
153 self.0.index_documents(&[document])?;
154 Ok(WriteReceipt { value: memory, revision })
155 }
156 pub fn get(&self, id: i64, filter: &ReadFilter) -> Result<Memory> {
157 let state = self.0.read()?;
158 storage::get(state.conn(), &RecordKey { id }, filter)
159 }
160 pub fn list(&self, request: &PageRequest) -> Result<Page<Memory>> {
161 storage::list(self.0.read()?.conn(), RecordKind::Memory, request)
162 }
163 pub fn delete(&self, id: i64, filter: &ReadFilter) -> Result<WriteReceipt<bool>> {
164 self.0.mutate(|tx| {
165 let key = RecordKey { id };
166 if !storage::matches_filter(tx, &key, filter)? { return Err(Error::NotFound(id.to_string())); }
167 storage::delete_record(tx, &key)
168 })
169 }
170 pub fn delete_by_filter(&self, filter: &ReadFilter) -> Result<WriteReceipt<usize>> {
173 self.0.mutate(|tx| {
174 let mut removed = 0;
175 for id in storage::select_ids(tx, filter, &[RecordKind::Memory])? {
176 if storage::delete_record(tx, &RecordKey { id })? { removed += 1; }
177 }
178 Ok(removed)
179 })
180 }
181 pub fn feedback(&self, request: &FeedbackRequest) -> Result<WriteReceipt<FeedbackReport>> {
182 request.policy.validate()?;
183 let recalled: BTreeSet<_> = request.recalled_ids.iter().copied().collect();
184 let useful: BTreeSet<_> = request.useful_ids.iter().copied().collect();
185 if !useful.is_subset(&recalled) { return Err(Error::Validation("useful_ids must be a subset of recalled_ids".into())); }
186 let receipt = self.0.mutate(|tx| {
187 let mut report = FeedbackReport { recalled: recalled.len(), ..Default::default() };
188 let mut documents = Vec::new();
189 let at = request.now_us.unwrap_or_else(storage::now_us);
190 for id in &recalled {
191 let key = RecordKey { id: *id };
192 let mut memory: Memory = storage::get(tx, &key, &request.filter)?;
193 if useful.contains(id) {
194 memory.state.strength = memory.state.strength.checked_add(request.policy.strength_boost).ok_or_else(|| Error::Validation("strength overflow".into()))?;
195 memory.state.useful_count = memory.state.useful_count.checked_add(1).ok_or_else(|| Error::Validation("useful_count overflow".into()))?;
196 memory.state.useful_score += request.policy.useful_score_boost;
197 memory.state.last_recalled_at_us = Some(at);
198 report.boosted += 1;
199 } else if !memory.state.pinned && request.policy.tier(memory.state.useful_score) == 1 {
200 memory.state.strength = (memory.state.strength - 1).max(0);
201 report.penalized += 1;
202 } else { continue; }
203 let (_, document) = upsert(tx, &as_input(&memory, at))?;
204 documents.push(document);
205 }
206 Ok((report, documents))
207 })?;
208 let WriteReceipt { value: (report, documents), revision } = receipt;
209 self.0.index_documents(&documents)?;
210 Ok(WriteReceipt { value: report, revision })
211 }
212 pub fn decay(&self, filter: &ReadFilter, policy: &DecayPolicy, at: Option<i64>) -> Result<WriteReceipt<DecayReport>> {
213 policy.validate()?;
214 let receipt = self.0.mutate(|tx| {
215 let at = at.unwrap_or_else(storage::now_us);
216 let keys = storage::select_keys(tx, filter, &[RecordKind::Memory], usize::MAX, None)?;
217 let cycle = i128::from(policy.tier0_cycle_days) * 86_400_000_000;
218 let mut documents = Vec::new();
219 let mut report = DecayReport::default();
220 for key in keys {
221 let mut memory: Memory = storage::get(tx, &key, filter)?;
222 if memory.state.pinned { continue; }
223 if policy.tier(memory.state.useful_score) == 0 && memory.state.strength > 0 {
224 let reference = memory.header.created_at_us.max(memory.state.last_recalled_at_us.unwrap_or(i64::MIN)).max(memory.state.last_decay_at_us.unwrap_or(i64::MIN));
225 let steps = (i128::from(at) - i128::from(reference)) / cycle;
226 if steps > 0 {
227 memory.state.strength = (i128::from(memory.state.strength) - steps).max(0) as i64;
228 memory.state.last_decay_at_us = Some((i128::from(reference) + steps * cycle) as i64);
229 let (_, document) = upsert(tx, &as_input(&memory, at))?;
230 documents.push(document);
231 report.decayed += 1;
232 }
233 }
234 if memory.state.strength == 0 && policy.tier(memory.state.useful_score) < 2 { report.retirement_candidates.push(key); }
235 }
236 Ok((report, documents))
237 })?;
238 let WriteReceipt { value: (report, documents), revision } = receipt;
239 self.0.index_documents(&documents)?;
240 Ok(WriteReceipt { value: report, revision })
241 }
242}