sim_lib_agent/memory/core/

advanced.rs

use sim_kernel::{CapabilityName, Cx, Error, Expr, Result, Symbol, Value};
use std::{
    cmp::Ordering,
    path::PathBuf,
    sync::{Arc, Mutex},
};

use crate::{FILE_WRITE_CAPABILITY, PersonaMemory, VectorMemory, cosine, embed};

use super::super::store::{
    append_memory_log, flatten_expr_text, load_memory_log, rewrite_memory_log, snapshot_entries,
    value_to_snapshot_expr,
};
use super::types::MemoryBackend;

#[derive(Clone)]
pub(super) struct VectorEntry {
    pub(super) expr: Expr,
    embedding: [f32; crate::embed::EMBED_DIM],
}

#[derive(Clone, Default)]
pub(super) struct PersonaState {
    profile: Vec<(Expr, Expr)>,
    notes: Vec<Expr>,
}

fn lock_vector_entries<'a>(
    entries: &'a Arc<Mutex<Vec<VectorEntry>>>,
) -> Result<std::sync::MutexGuard<'a, Vec<VectorEntry>>> {
    entries
        .lock()
        .map_err(|_| Error::PoisonedLock("vector memory entries"))
}

fn lock_persona_state<'a>(
    state: &'a Arc<Mutex<PersonaState>>,
) -> Result<std::sync::MutexGuard<'a, PersonaState>> {
    state
        .lock()
        .map_err(|_| Error::PoisonedLock("persona memory state"))
}

impl VectorMemory {
    pub(crate) fn open(path: Option<PathBuf>, codecs: Vec<Symbol>) -> Result<Self> {
        let entries = match &path {
            Some(path) => load_memory_log(path)?
                .into_iter()
                .map(VectorEntry::from_expr)
                .collect(),
            None => Vec::new(),
        };
        Ok(Self {
            symbol: Symbol::qualified("memory", "vector"),
            capabilities: path
                .as_ref()
                .map(|_| vec![CapabilityName::new(FILE_WRITE_CAPABILITY)])
                .unwrap_or_default(),
            address: sim_lib_server::ServerAddress::Local,
            codecs,
            path,
            entries: Arc::new(Mutex::new(entries)),
        })
    }
}

impl PersonaMemory {
    pub(crate) fn open(path: Option<PathBuf>, codecs: Vec<Symbol>) -> Result<Self> {
        let state = match &path {
            Some(path) => {
                let entries = load_memory_log(path)?;
                match entries.as_slice() {
                    [] => PersonaState::default(),
                    [snapshot] => persona_state_from_snapshot(snapshot.clone())?,
                    _ => {
                        return Err(Error::HostError(format!(
                            "persona memory {} expected at most one snapshot entry",
                            path.display()
                        )));
                    }
                }
            }
            None => PersonaState::default(),
        };
        Ok(Self {
            symbol: Symbol::qualified("memory", "persona"),
            capabilities: path
                .as_ref()
                .map(|_| vec![CapabilityName::new(FILE_WRITE_CAPABILITY)])
                .unwrap_or_default(),
            address: sim_lib_server::ServerAddress::Local,
            codecs,
            path,
            state: Arc::new(Mutex::new(state)),
        })
    }
}

impl VectorEntry {
    fn from_expr(expr: Expr) -> Self {
        let embedding = embed(&flatten_expr_text(&expr));
        Self { expr, embedding }
    }
}

impl PersonaState {
    fn snapshot(&self) -> Expr {
        Expr::List(vec![Expr::Map(vec![
            (
                Expr::Symbol(Symbol::new("state")),
                Expr::Map(self.profile.clone()),
            ),
            (
                Expr::Symbol(Symbol::new("notes")),
                Expr::List(self.notes.clone()),
            ),
        ])])
    }
}

fn persona_state_from_snapshot(snapshot: Expr) -> Result<PersonaState> {
    let (Expr::List(items) | Expr::Vector(items)) = snapshot else {
        return Err(Error::TypeMismatch {
            expected: "persona snapshot list",
            found: "non-list",
        });
    };
    let Some(entry) = items.first() else {
        return Ok(PersonaState::default());
    };
    let Expr::Map(entries) = entry else {
        return Err(Error::TypeMismatch {
            expected: "persona snapshot map",
            found: "non-map",
        });
    };
    let mut profile = Vec::new();
    let mut notes = Vec::new();
    for (key, value) in entries {
        match key {
            Expr::Symbol(symbol) if symbol.name.as_ref() == "state" => {
                let Expr::Map(state_entries) = value else {
                    return Err(Error::TypeMismatch {
                        expected: "persona state map",
                        found: "non-map",
                    });
                };
                profile = state_entries.clone();
            }
            Expr::Symbol(symbol) if symbol.name.as_ref() == "notes" => {
                notes = snapshot_entries(value.clone())?;
            }
            _ => {}
        }
    }
    Ok(PersonaState { profile, notes })
}

impl MemoryBackend for VectorMemory {
    fn append(&self, cx: &mut Cx, msg: Value) -> Result<()> {
        if self.path.is_some() {
            cx.require(&CapabilityName::new(FILE_WRITE_CAPABILITY))?;
        }
        let expr = value_to_snapshot_expr(cx, msg)?;
        let entry = VectorEntry::from_expr(expr.clone());
        lock_vector_entries(&self.entries)?.push(entry);
        if let Some(path) = &self.path {
            append_memory_log(path, &expr)?;
        }
        Ok(())
    }

    fn recent(&self, cx: &mut Cx, count: u32) -> Result<Vec<Value>> {
        let entries = lock_vector_entries(&self.entries)?;
        let count = usize::try_from(count).unwrap_or(usize::MAX);
        entries[entries.len().saturating_sub(count)..]
            .iter()
            .map(|entry| crate::expr_to_value(cx, &entry.expr))
            .collect()
    }

    fn search(&self, cx: &mut Cx, query: Expr, k: u32) -> Result<Vec<Value>> {
        let query_embedding = embed(&flatten_expr_text(&query));
        let mut scored = lock_vector_entries(&self.entries)?
            .iter()
            .enumerate()
            .map(|(index, entry)| {
                (
                    index,
                    cosine(&query_embedding, &entry.embedding),
                    entry.expr.clone(),
                )
            })
            .collect::<Vec<_>>();
        scored.sort_by(|left, right| {
            right
                .1
                .partial_cmp(&left.1)
                .unwrap_or(Ordering::Equal)
                .then(left.0.cmp(&right.0))
        });
        scored
            .into_iter()
            .take(usize::try_from(k).unwrap_or(usize::MAX))
            .map(|(_, _, expr)| crate::expr_to_value(cx, &expr))
            .collect()
    }

    fn snapshot(&self, _cx: &mut Cx) -> Result<Expr> {
        Ok(Expr::List(
            lock_vector_entries(&self.entries)?
                .iter()
                .map(|entry| entry.expr.clone())
                .collect(),
        ))
    }

    fn restore(&self, cx: &mut Cx, snap: Expr) -> Result<()> {
        if self.path.is_some() {
            cx.require(&CapabilityName::new(FILE_WRITE_CAPABILITY))?;
        }
        let entries = snapshot_entries(snap)?
            .into_iter()
            .map(VectorEntry::from_expr)
            .collect::<Vec<_>>();
        if let Some(path) = &self.path {
            let exprs = entries
                .iter()
                .map(|entry| entry.expr.clone())
                .collect::<Vec<_>>();
            rewrite_memory_log(path, &exprs)?;
        }
        *lock_vector_entries(&self.entries)? = entries;
        Ok(())
    }
}

impl MemoryBackend for PersonaMemory {
    fn append(&self, cx: &mut Cx, msg: Value) -> Result<()> {
        if self.path.is_some() {
            cx.require(&CapabilityName::new(FILE_WRITE_CAPABILITY))?;
        }
        let note = value_to_snapshot_expr(cx, msg)?;
        let mut state = lock_persona_state(&self.state)?;
        state.notes.push(note);
        if let Some(path) = &self.path {
            rewrite_memory_log(path, &[state.snapshot()])?;
        }
        Ok(())
    }

    fn recent(&self, cx: &mut Cx, count: u32) -> Result<Vec<Value>> {
        let state = lock_persona_state(&self.state)?;
        let count = usize::try_from(count).unwrap_or(usize::MAX);
        state.notes[state.notes.len().saturating_sub(count)..]
            .iter()
            .map(|entry| crate::expr_to_value(cx, entry))
            .collect()
    }

    fn search(&self, cx: &mut Cx, query: Expr, k: u32) -> Result<Vec<Value>> {
        let query_terms = flatten_expr_text(&query)
            .split_whitespace()
            .map(str::to_owned)
            .collect::<Vec<_>>();
        let state = lock_persona_state(&self.state)?;
        let filtered = state
            .notes
            .iter()
            .rev()
            .filter(|expr| {
                let haystack = flatten_expr_text(expr);
                query_terms.iter().all(|term| haystack.contains(term))
            })
            .take(usize::try_from(k).unwrap_or(usize::MAX))
            .cloned()
            .collect::<Vec<_>>();
        filtered
            .into_iter()
            .rev()
            .map(|expr| crate::expr_to_value(cx, &expr))
            .collect()
    }

    fn snapshot(&self, _cx: &mut Cx) -> Result<Expr> {
        Ok(lock_persona_state(&self.state)?.snapshot())
    }

    fn restore(&self, cx: &mut Cx, snap: Expr) -> Result<()> {
        if self.path.is_some() {
            cx.require(&CapabilityName::new(FILE_WRITE_CAPABILITY))?;
        }
        let state = persona_state_from_snapshot(snap)?;
        if let Some(path) = &self.path {
            rewrite_memory_log(path, &[state.snapshot()])?;
        }
        *lock_persona_state(&self.state)? = state;
        Ok(())
    }
}