use crate::screening::{self, RejectReason, ScreeningVerdict};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryRecord {
pub id: String,
pub scope: String,
pub content: String,
pub digest: String,
pub subject_key: String,
pub created_at: String,
pub access_count: u32,
pub pinned: bool,
}
#[derive(Debug, Clone)]
pub struct NewMemory {
pub id: String,
pub scope: String,
pub content: String,
pub digest: String,
pub subject_key: String,
pub created_at: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "reason")]
pub enum WriteRejection {
Screened { detail: RejectReason },
DuplicateId,
}
impl From<RejectReason> for WriteRejection {
fn from(detail: RejectReason) -> Self {
WriteRejection::Screened { detail }
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MemoryStore {
records: std::collections::BTreeMap<String, MemoryRecord>,
}
impl MemoryStore {
pub fn new() -> Self {
Self::default()
}
pub fn write(&mut self, candidate: NewMemory) -> Result<&MemoryRecord, WriteRejection> {
if self.records.contains_key(&candidate.id) {
return Err(WriteRejection::DuplicateId);
}
match screening::screen(&candidate.content) {
ScreeningVerdict::Reject { reason } => Err(WriteRejection::from(reason)),
ScreeningVerdict::Admit => {
let record = MemoryRecord {
id: candidate.id.clone(),
scope: candidate.scope,
content: candidate.content,
digest: candidate.digest,
subject_key: candidate.subject_key,
created_at: candidate.created_at,
access_count: 0,
pinned: false,
};
self.records.insert(candidate.id.clone(), record);
Ok(self.records.get(&candidate.id).expect("just inserted"))
}
}
}
pub fn delete(&mut self, id: &str) -> Option<MemoryRecord> {
self.records.remove(id)
}
pub fn get(&self, id: &str) -> Option<&MemoryRecord> {
self.records.get(id)
}
pub fn set_pinned(&mut self, id: &str, pinned: bool) -> bool {
if let Some(record) = self.records.get_mut(id) {
record.pinned = pinned;
true
} else {
false
}
}
pub fn record_access(&mut self, id: &str) {
if let Some(record) = self.records.get_mut(id) {
record.access_count = record.access_count.saturating_add(1);
}
}
pub fn list(&self, scope: Option<&str>) -> Vec<&MemoryRecord> {
self.records
.values()
.filter(|record| match scope {
Some(s) => record.scope == s,
None => true,
})
.collect()
}
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
}