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 self.0.mutate(|tx| upsert(tx, &input).map(|(memory, document)| (memory, vec![document])))
122 }
123 pub fn upsert_many(&self, inputs: &[MemoryInput]) -> Result<WriteReceipt<Vec<Memory>>> {
124 self.0.mutate(|tx| {
125 let pairs = inputs.iter().map(|input| upsert(tx, input)).collect::<Result<Vec<_>>>()?;
126 Ok(pairs.into_iter().unzip())
127 })
128 }
129 pub fn upsert_by_judgment(&self, mut input: MemoryInput) -> Result<WriteReceipt<Memory>> {
130 self.0.mutate(|tx| {
131 let mut stmt = tx.prepare("SELECT id FROM records WHERE namespace_id=(SELECT id FROM strings WHERE text=?1)
132 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")?;
133 let ids = stmt.query_map(params![text::normalized_tag(&input.record.namespace), text::normalized_tag(&input.record.scope),
134 RecordKind::Memory.code(), text::normalized_tag(&input.judgment)], |r| r.get::<_, i64>(0))?.collect::<std::result::Result<Vec<_>, _>>()?;
135 if ids.len() > 1 { return Err(Error::Conflict("multiple memories have this judgment; update by ID".into())); }
136 if let Some(id) = ids.first() {
137 if input.record.id.is_some_and(|given| given != *id) { return Err(Error::Conflict("judgment belongs to a different ID".into())); }
138 input.record.id = Some(*id);
139 let key = RecordKey { id: *id };
140 let existing: Memory = serde_json::from_value(storage::record_value(tx, &key)?.ok_or_else(|| Error::NotFound(id.to_string()))?)?;
141 let mut metadata = existing.header.metadata;
142 metadata.extend(input.record.metadata.clone());
143 input.record.metadata = metadata;
144 if input.record.evidence.is_empty() { input.record.evidence = existing.header.evidence; }
145 }
146 upsert(tx, &input).map(|(memory, document)| (memory, vec![document]))
147 })
148 }
149 pub fn get(&self, id: i64, filter: &ReadFilter) -> Result<Memory> {
150 let state = self.0.read()?;
151 storage::get(state.conn(), &RecordKey { id }, filter)
152 }
153 pub fn list(&self, request: &PageRequest) -> Result<Page<Memory>> {
154 storage::list(self.0.read()?.conn(), RecordKind::Memory, request)
155 }
156 pub fn delete(&self, ids: &[i64], filter: &ReadFilter) -> Result<WriteReceipt<usize>> {
160 let removed = self.0.delete_flow(ids, filter, RecordKind::Memory, |_, _| Ok(Vec::new()))?;
161 let revision = storage::current_revision(self.0.read()?.conn())?;
162 Ok(WriteReceipt { value: removed, revision })
163 }
164 pub fn feedback(&self, request: &FeedbackRequest) -> Result<WriteReceipt<FeedbackReport>> {
165 request.policy.validate()?;
166 let recalled: BTreeSet<_> = request.recalled_ids.iter().copied().collect();
167 let useful: BTreeSet<_> = request.useful_ids.iter().copied().collect();
168 if !useful.is_subset(&recalled) { return Err(Error::Validation("useful_ids must be a subset of recalled_ids".into())); }
169 self.0.mutate(|tx| {
170 let mut report = FeedbackReport { recalled: recalled.len(), ..Default::default() };
171 let mut documents = Vec::new();
172 let at = request.now_us.unwrap_or_else(storage::now_us);
173 for id in &recalled {
174 let key = RecordKey { id: *id };
175 let mut memory: Memory = storage::get(tx, &key, &request.filter)?;
176 if useful.contains(id) {
177 memory.state.strength = memory.state.strength.checked_add(request.policy.strength_boost).ok_or_else(|| Error::Validation("strength overflow".into()))?;
178 memory.state.useful_count = memory.state.useful_count.checked_add(1).ok_or_else(|| Error::Validation("useful_count overflow".into()))?;
179 memory.state.useful_score += request.policy.useful_score_boost;
180 memory.state.last_recalled_at_us = Some(at);
181 report.boosted += 1;
182 } else if !memory.state.pinned && request.policy.tier(memory.state.useful_score) == 1 {
183 memory.state.strength = (memory.state.strength - 1).max(0);
184 report.penalized += 1;
185 } else { continue; }
186 let (_, document) = upsert(tx, &as_input(&memory, at))?;
187 documents.push(document);
188 }
189 Ok((report, documents))
190 })
191 }
192 pub fn decay(&self, filter: &ReadFilter, policy: &DecayPolicy, at: Option<i64>) -> Result<WriteReceipt<DecayReport>> {
193 policy.validate()?;
194 self.0.mutate(|tx| {
195 let at = at.unwrap_or_else(storage::now_us);
196 let keys = storage::select_keys(tx, filter, &[RecordKind::Memory], usize::MAX, None)?;
197 let cycle = i128::from(policy.tier0_cycle_days) * 86_400_000_000;
198 let mut documents = Vec::new();
199 let mut report = DecayReport::default();
200 for key in keys {
201 let mut memory: Memory = storage::get(tx, &key, filter)?;
202 if memory.state.pinned { continue; }
203 if policy.tier(memory.state.useful_score) == 0 && memory.state.strength > 0 {
204 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));
205 let steps = (i128::from(at) - i128::from(reference)) / cycle;
206 if steps > 0 {
207 memory.state.strength = (i128::from(memory.state.strength) - steps).max(0) as i64;
208 memory.state.last_decay_at_us = Some((i128::from(reference) + steps * cycle) as i64);
209 let (_, document) = upsert(tx, &as_input(&memory, at))?;
210 documents.push(document);
211 report.decayed += 1;
212 }
213 }
214 if memory.state.strength == 0 && policy.tier(memory.state.useful_score) < 2 { report.retirement_candidates.push(key); }
215 }
216 Ok((report, documents))
217 })
218 }
219}