Skip to main content

p_memory/
memory.rs

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    /// None preserves an existing memory's lifecycle state.
38    #[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        // 入库走一次写锁(见 `mutate`);索引文档在写锁外交给索引 writer。
120        // 索引提交交给 `update_index`,向量化交给使用方调 `embeddings().sync`。
121        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    /// 批量删除记忆:主库查 id,命中才继续;单条也是批量的一种。
157    /// 写锁内落主库标记、删向量行与主库行,出锁后摘索引词条;
158    /// 断电留下「正在删除」标记的,下次开机把这条删除做完。
159    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}