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        let receipt = self.0.mutate(|tx| upsert(tx, &input))?;
120        let WriteReceipt { value: (memory, document), revision } = receipt;
121        // 写入只做两件事:入库(上面的 mutate)与把索引文档交给 writer。
122        // 索引提交交给 `update_index`,向量化交给 `embeddings().sync` 或库内线程。
123        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    /// 删除过滤条件命中的全部记忆,返回删除条数。空命中返回 0。
171    /// 记忆没有子表,删一条与删一批走的是同一个 `delete_record`。
172    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}