use crate::server::protocol::RemoteRole;
use serde::{Deserialize, Serialize};
const PROBATION_DURATION_MS: i64 = 7 * 24 * 60 * 60 * 1_000;
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct MemoryKey {
pub id: i64,
pub version: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub namespace: Option<String>,
}
impl MemoryKey {
pub const fn local(id: i64, version: u64) -> Self {
Self {
id,
version,
namespace: None,
}
}
pub fn remote(namespace: String, id: i64, version: u64) -> Self {
Self {
id,
version,
namespace: Some(namespace),
}
}
pub const fn is_local(&self) -> bool {
self.namespace.is_none()
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MemoryRecord {
pub key: MemoryKey,
pub content: String,
pub created_at_ms: i64,
pub updated_at_ms: i64,
pub last_scanned_at_ms: Option<i64>,
pub scan_count: u64,
pub last_used_at_ms: Option<i64>,
pub use_count: u64,
pub probation_until_ms: Option<i64>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct MemoryCandidate {
pub key: MemoryKey,
pub preview: String,
pub score: f64,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct MemoryScan {
pub abstained: bool,
pub candidates: Vec<MemoryCandidate>,
}
impl MemoryScan {
pub fn rank(query: &str, memories: &[MemoryRecord], limit: usize) -> Self {
let candidates = crate::retrieval::rank(query, memories, limit);
Self {
abstained: candidates.is_empty(),
candidates,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MemoryLimits {
pub content_bytes: usize,
pub records: usize,
pub total_content_bytes: usize,
pub database_bytes: usize,
pub scan_results: usize,
pub query_bytes: usize,
pub probation_duration_ms: i64,
}
impl MemoryLimits {
pub const PRODUCTION: Self = Self {
content_bytes: 1_024,
records: 512,
total_content_bytes: 256 * 1_024,
database_bytes: 4 * 1_024 * 1_024,
scan_results: 5,
query_bytes: 512,
probation_duration_ms: PROBATION_DURATION_MS,
};
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum MemorySource {
Local,
Remote,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct MemoryAccess {
pub source: MemorySource,
pub namespace: Option<String>,
pub role: Option<RemoteRole>,
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub struct MemoryImportReport {
pub inserted: usize,
pub skipped: usize,
}
#[derive(Clone, Debug)]
pub(crate) struct StoredMemory {
pub(crate) namespace: Option<String>,
pub(crate) id: i64,
pub(crate) content: String,
pub(crate) created_at_ms: i64,
pub(crate) updated_at_ms: i64,
pub(crate) last_scanned_at_ms: Option<i64>,
pub(crate) scan_count: u64,
pub(crate) last_used_at_ms: Option<i64>,
pub(crate) use_count: u64,
pub(crate) probation_until_ms: Option<i64>,
pub(crate) version: u64,
}
impl StoredMemory {
pub(crate) fn key(&self) -> MemoryKey {
match &self.namespace {
Some(namespace) => MemoryKey::remote(namespace.clone(), self.id, self.version),
None => MemoryKey::local(self.id, self.version),
}
}
}
impl From<StoredMemory> for MemoryRecord {
fn from(memory: StoredMemory) -> Self {
Self {
key: memory.key(),
content: memory.content,
created_at_ms: memory.created_at_ms,
updated_at_ms: memory.updated_at_ms,
last_scanned_at_ms: memory.last_scanned_at_ms,
scan_count: memory.scan_count,
last_used_at_ms: memory.last_used_at_ms,
use_count: memory.use_count,
probation_until_ms: memory.probation_until_ms,
}
}
}
pub fn normalize_identity(content: &str) -> String {
content
.split_whitespace()
.map(str::to_lowercase)
.collect::<Vec<_>>()
.join(" ")
}