use crate::agent::Laser;
use crate::error::LaserError;
use crate::provenance::{AgentTopic, Provenance};
use crate::types::{AgentId, ConversationId, IdError};
use iggy::prelude::*;
use laser_wire::graph::SourceRef;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::fmt::{self, Display, Formatter};
use std::str::FromStr;
use tokio::sync::Mutex;
use ulid::Ulid;
const READ_BATCH: u32 = 1000;
const CONSOLIDATION_WINDOW: usize = 10_000;
pub const DEFAULT_MEMORY_TOPIC_TTL: std::time::Duration =
std::time::Duration::from_secs(30 * 24 * 60 * 60);
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryKind {
#[default]
Fact,
Message,
Summary,
Entity,
Feedback,
Procedure,
}
impl MemoryKind {
pub const fn code(self) -> u8 {
match self {
MemoryKind::Fact => 1,
MemoryKind::Message => 2,
MemoryKind::Summary => 3,
MemoryKind::Entity => 4,
MemoryKind::Feedback => 5,
MemoryKind::Procedure => 6,
}
}
pub fn from_word(word: &str) -> MemoryKind {
match word {
"message" => MemoryKind::Message,
"summary" => MemoryKind::Summary,
"entity" => MemoryKind::Entity,
"feedback" => MemoryKind::Feedback,
"procedure" => MemoryKind::Procedure,
_ => MemoryKind::Fact,
}
}
pub const fn class(self) -> MemoryClass {
match self {
MemoryKind::Message => MemoryClass::Episodic,
MemoryKind::Procedure => MemoryClass::Procedural,
MemoryKind::Fact | MemoryKind::Summary | MemoryKind::Entity | MemoryKind::Feedback => {
MemoryClass::Semantic
}
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryClass {
Episodic,
Semantic,
Procedural,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Lifetime {
#[default]
Session,
Durable,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RecallStrategy {
#[default]
Auto,
Recent,
Semantic,
Keyword,
Graph,
Temporal,
Hybrid,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MemoryId(Ulid);
impl MemoryId {
pub fn new() -> Self {
Self(Ulid::generate())
}
pub fn as_u128(self) -> u128 {
self.0.0
}
pub fn from_u128(value: u128) -> Self {
Self(Ulid(value))
}
pub fn content(owner: &MemoryScope, kind: MemoryKind, body: &[u8]) -> Self {
let stream = owner.stream.as_deref().unwrap_or("").as_bytes();
let agent = owner.agent.as_ref().map(AgentId::as_str).unwrap_or("");
let kind = [kind.code()];
let segments: [&[u8]; 5] = [stream, &[0], agent.as_bytes(), &[0], &kind];
let mut all = segments.to_vec();
all.push(body);
Self::from_u128(laser_wire::hashing::content_id(&all))
}
}
impl Serialize for MemoryId {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(&self.to_string())
}
}
impl<'de> Deserialize<'de> for MemoryId {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let text = String::deserialize(deserializer)?;
text.parse().map_err(serde::de::Error::custom)
}
}
impl Default for MemoryId {
fn default() -> Self {
Self::new()
}
}
impl Display for MemoryId {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
Display::fmt(&self.0, f)
}
}
impl FromStr for MemoryId {
type Err = IdError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ulid::from_string(s)
.map(Self)
.map_err(|_| IdError::InvalidUlid(s.to_owned()))
}
}
#[derive(Debug, Clone, Default, bon::Builder)]
pub struct MemoryScope {
pub stream: Option<String>,
pub user: Option<String>,
pub agent: Option<AgentId>,
pub conversation: Option<ConversationId>,
pub app: Option<String>,
#[builder(default)]
pub lifetime: Lifetime,
}
#[derive(Debug, Clone)]
pub struct MemoryItem {
pub id: MemoryId,
pub payload: Vec<u8>,
pub provenance: Provenance,
pub kind: MemoryKind,
pub score: Option<f32>,
pub signals: Vec<RecallSignal>,
pub source: Option<SourceRef>,
}
#[cfg(test)]
impl MemoryItem {
fn plain(id: MemoryId, payload: Vec<u8>, provenance: Provenance) -> Self {
Self {
id,
payload,
provenance,
kind: MemoryKind::Fact,
score: None,
signals: Vec::new(),
source: None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct RecallSignal {
pub strategy: RecallStrategy,
pub rank: usize,
pub score: Option<f32>,
}
pub fn to_context_block(items: &[MemoryItem], token_budget: Option<usize>) -> String {
let mut block = String::new();
let mut spent = 0usize;
let mut rendered = 0usize;
for item in items {
let text = String::from_utf8_lossy(&item.payload);
let cost = estimate_tokens(&text);
if let Some(budget) = token_budget
&& rendered > 0
&& spent + cost > budget
{
break;
}
if rendered > 0 {
block.push_str("\n\n");
}
block.push_str(&text);
spent += cost;
rendered += 1;
}
let omitted = items.len() - rendered;
if omitted > 0 {
if rendered > 0 {
block.push_str("\n\n");
}
block.push_str(&format!(
"[... {omitted} more recalled item(s) omitted ...]"
));
}
block
}
fn estimate_tokens(text: &str) -> usize {
text.len().div_ceil(4)
}
#[derive(Debug, Clone, bon::Builder)]
pub struct MemoryQuery {
#[builder(default = 50)]
pub limit: usize,
pub token_budget: Option<usize>,
pub agent: Option<AgentId>,
pub semantic: Option<String>,
#[builder(default)]
pub strategy: RecallStrategy,
}
#[derive(Debug, Clone)]
pub struct Feedback {
pub target: MemoryId,
pub weight: f32,
pub note: Option<String>,
}
impl Feedback {
pub fn new(target: MemoryId, weight: f32) -> Self {
Self {
target,
weight,
note: None,
}
}
}
#[trait_variant::make(Memory: Send)]
pub trait LocalMemory {
async fn remember(&self, scope: &MemoryScope, payload: Vec<u8>)
-> Result<MemoryId, LaserError>;
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError>;
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError>;
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError>;
}
pub trait DynMemory: Send + Sync {
fn remember<'a>(
&'a self,
scope: &'a MemoryScope,
payload: Vec<u8>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
>;
fn recall<'a>(
&'a self,
scope: &'a MemoryScope,
query: &'a MemoryQuery,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>,
>;
fn improve<'a>(
&'a self,
scope: &'a MemoryScope,
feedback: Feedback,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
>;
fn forget<'a>(
&'a self,
scope: &'a MemoryScope,
id: MemoryId,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), LaserError>> + Send + 'a>>;
}
impl<M: Memory + Send + Sync> DynMemory for M {
fn remember<'a>(
&'a self,
scope: &'a MemoryScope,
payload: Vec<u8>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
> {
Box::pin(Memory::remember(self, scope, payload))
}
fn recall<'a>(
&'a self,
scope: &'a MemoryScope,
query: &'a MemoryQuery,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>,
> {
Box::pin(Memory::recall(self, scope, query))
}
fn improve<'a>(
&'a self,
scope: &'a MemoryScope,
feedback: Feedback,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
> {
Box::pin(Memory::improve(self, scope, feedback))
}
fn forget<'a>(
&'a self,
scope: &'a MemoryScope,
id: MemoryId,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), LaserError>> + Send + 'a>>
{
Box::pin(Memory::forget(self, scope, id))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
enum MemoryLogEntry {
Item {
id: MemoryId,
kind: MemoryKind,
body: Vec<u8>,
},
Forget { target: MemoryId },
Feedback { target: MemoryId, weight: f32 },
}
impl MemoryLogEntry {
fn encode(&self) -> Result<Vec<u8>, LaserError> {
laser_wire::framing::encode_named(self)
.map_err(|error| LaserError::Codec(format!("encode memory entry: {error}")))
}
fn decode(payload: &[u8]) -> Result<Self, LaserError> {
laser_wire::framing::decode_named(payload)
.map_err(|error| LaserError::Codec(format!("decode memory entry: {error}")))
}
}
pub struct LogMemory {
laser: Laser,
stream: Option<String>,
topic: Identifier,
namespace: String,
projection: Mutex<Projection>,
}
#[derive(Default)]
struct Projection {
items: Vec<MemoryItem>,
forgotten: HashSet<MemoryId>,
feedback: std::collections::HashMap<MemoryId, f32>,
named: std::collections::HashMap<String, Vec<u8>>,
offsets: Vec<u64>,
}
impl Projection {
fn absorb(&mut self, payload: &[u8], provenance: Provenance) {
let Ok(entry) = laser_wire::framing::decode_named::<MemoryLogEntry>(payload) else {
self.absorb_named(payload);
return;
};
match entry {
MemoryLogEntry::Forget { target } => {
self.items.retain(|item| item.id != target);
self.forgotten.insert(target);
}
MemoryLogEntry::Feedback { target, weight } => {
*self.feedback.entry(target).or_insert(0.0) += weight;
}
MemoryLogEntry::Item { id, kind, body } => {
if self.forgotten.contains(&id) {
return;
}
if self.items.iter().any(|item| item.id == id) {
return;
}
self.items.push(MemoryItem {
id,
payload: body,
provenance,
source: None,
kind,
score: None,
signals: Vec::new(),
});
}
}
}
fn absorb_named(&mut self, payload: &[u8]) {
let Ok(record) =
laser_wire::framing::decode_named::<laser_wire::memory::MemoryRecord>(payload)
else {
return;
};
match record {
laser_wire::memory::MemoryRecord::Item { id, body, .. } => {
self.named.insert(id, body);
}
laser_wire::memory::MemoryRecord::Forget { target } => {
self.named.remove(&target);
}
laser_wire::memory::MemoryRecord::Feedback { .. } => {}
}
}
}
impl LogMemory {
pub fn new(laser: Laser) -> Self {
Self {
laser,
stream: None,
topic: AgentTopic::Audit.as_identifier(),
namespace: AgentTopic::Audit.name().unwrap_or("memory").to_owned(),
projection: Mutex::new(Projection::default()),
}
}
pub fn in_namespace(laser: Laser, namespace: impl Into<String>) -> Self {
Self {
laser,
stream: None,
topic: AgentTopic::Audit.as_identifier(),
namespace: namespace.into(),
projection: Mutex::new(Projection::default()),
}
}
pub fn on_topic(laser: Laser, topic: Identifier) -> Self {
Self::on_stream_topic(laser, None, topic)
}
pub fn on_stream_topic(laser: Laser, stream: Option<String>, topic: Identifier) -> Self {
let namespace = topic.get_string_value().unwrap_or_default();
Self {
laser,
stream,
topic,
namespace,
projection: Mutex::new(Projection::default()),
}
}
pub fn on_topic_named(laser: Laser, topic: &str) -> Result<Self, LaserError> {
Ok(Self::on_topic(laser, Identifier::named(topic)?))
}
pub fn on_stream_topic_named(
laser: Laser,
stream: Option<String>,
topic: &str,
) -> Result<Self, LaserError> {
Ok(Self::on_stream_topic(
laser,
stream,
Identifier::named(topic)?,
))
}
async fn govern_payload(
&self,
provenance: &Provenance,
payload: Vec<u8>,
) -> Result<Vec<u8>, LaserError> {
let topic = self.topic.to_string();
let stream = match &self.stream {
Some(stream) => stream.as_str(),
None => self.laser.stream_required()?,
};
let action = crate::govern::GovernedAction {
kind: crate::govern::ActionKind::MemoryWrite,
stream,
topic: &topic,
source: provenance.agent.as_ref().map(|agent| agent.as_str()),
target: None,
conversation: Some(provenance.conversation_id),
correlation: None,
operation: None,
tool: None,
on_behalf_of: None,
purpose: None,
data_classification: None,
payload: &payload,
signed: false,
counters: crate::govern::ActionCounters::default(),
};
Ok(self.laser.govern(action).await?.unwrap_or(payload))
}
async fn publish_unchecked(
&self,
provenance: &Provenance,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<(), LaserError> {
let mut headers: std::collections::BTreeMap<HeaderKey, HeaderValue> =
provenance.try_into()?;
put_header(
&mut headers,
laser_wire::headers::MEMORY_NAMESPACE,
&self.namespace,
)?;
if let Some(user) = &scope.user {
put_header(&mut headers, laser_wire::headers::MEMORY_USER, user)?;
}
if let Some(app) = &scope.app {
put_header(&mut headers, laser_wire::headers::MEMORY_APP, app)?;
}
let key = provenance.partition_key();
let topic = self.topic.to_string();
match &self.stream {
Some(stream) => self
.laser
.send_with_headers_on(stream, &topic, payload, headers, Some(&key))
.await
.map(|_| ()),
None => self
.laser
.send_with_headers(&topic, payload, headers, Some(&key))
.await
.map(|_| ()),
}
}
async fn publish(
&self,
provenance: &Provenance,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<(), LaserError> {
let payload = self.govern_payload(provenance, payload).await?;
self.publish_unchecked(provenance, scope, payload).await
}
fn provenance(scope: &MemoryScope, idempotency_key: String) -> Provenance {
let mut provenance = Provenance::builder()
.conversation_id(scope.conversation.unwrap_or_default())
.build();
provenance.agent = scope.agent.clone();
provenance.idempotency_key = Some(idempotency_key);
provenance
}
fn named_provenance(id: String) -> Provenance {
let mut provenance = Provenance::builder()
.conversation_id(ConversationId::derive(&id))
.build();
provenance.idempotency_key = Some(id);
provenance
}
pub(crate) async fn append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
let provenance = Self::provenance(scope, id.to_string());
let body = self.govern_payload(&provenance, payload).await?;
let entry = MemoryLogEntry::Item { id, kind, body };
self.publish_unchecked(&provenance, scope, entry.encode()?)
.await?;
Ok(id)
}
fn named_key(&self, key: &str) -> String {
format!("{}/{}", self.namespace, key)
}
pub async fn set_named(&self, key: &str, body: Vec<u8>) -> Result<(), LaserError> {
let id = self.named_key(key);
let record = laser_wire::memory::MemoryRecord::Item {
id: id.clone(),
kind: "fact".to_owned(),
body: body.clone(),
};
let payload = laser_wire::framing::encode_named(&record)
.map_err(|error| LaserError::Codec(format!("encode named memory item: {error}")))?;
let scope = MemoryScope::default();
let provenance = Self::named_provenance(id.clone());
self.publish(&provenance, &scope, payload).await?;
self.projection.lock().await.named.insert(id, body);
Ok(())
}
pub async fn update_named(&self, key: &str, patch: Vec<u8>) -> Result<(), LaserError> {
let current = self.fetch_named_folded(key).await?;
let merged = merge_named(current.as_deref(), &patch)?;
self.set_named(key, merged).await
}
pub async fn forget_named(&self, key: &str) -> Result<(), LaserError> {
let id = self.named_key(key);
let record = laser_wire::memory::MemoryRecord::Forget { target: id.clone() };
let payload = laser_wire::framing::encode_named(&record)
.map_err(|error| LaserError::Codec(format!("encode named memory forget: {error}")))?;
let scope = MemoryScope::default();
let provenance = Self::named_provenance(id.clone());
self.publish(&provenance, &scope, payload).await?;
self.projection.lock().await.named.remove(&id);
Ok(())
}
pub async fn fetch_named(&self, key: &str) -> Result<Option<Vec<u8>>, LaserError> {
#[cfg(feature = "kv")]
{
self.laser
.kv(self.namespace.clone())
.get(self.named_key(key))
.await
}
#[cfg(not(feature = "kv"))]
{
let _ = key;
Err(LaserError::unsupported(
"memory fetch",
"named point-state read reads the managed key-value view. Build with the `kv` \
feature, or call `fetch_named_folded` to fold the topic in process",
))
}
}
pub async fn fetch_named_folded(&self, key: &str) -> Result<Option<Vec<u8>>, LaserError> {
self.catch_up().await?;
let id = self.named_key(key);
Ok(self.projection.lock().await.named.get(&id).cloned())
}
async fn catch_up(&self) -> Result<(), LaserError> {
let stream = match &self.stream {
Some(stream) => Identifier::named(stream)?,
None => Identifier::named(self.laser.stream_required()?)?,
};
let topic = self.topic.clone();
let Some(details) = self.laser.client().get_topic(&stream, &topic).await? else {
return Ok(());
};
let partitions = crate::poll::bounded_partitions(details.partitions_count);
let consumer = Consumer::new(Identifier::named("laser-log-memory")?);
let from = {
let mut projection = self.projection.lock().await;
if projection.offsets.len() < partitions as usize {
projection.offsets.resize(partitions as usize, 0);
}
projection.offsets.clone()
};
let mut drained: Vec<(u32, u64, Vec<IggyMessage>)> = Vec::new();
for partition in 0..partitions {
let batch = crate::poll::drain_partition(
self.laser.client(),
&stream,
&topic,
&consumer,
partition,
from[partition as usize],
READ_BATCH,
)
.await?;
if !batch.messages.is_empty() {
drained.push((partition, batch.next_offset, batch.messages));
}
}
if drained.is_empty() {
return Ok(());
}
let mut projection = self.projection.lock().await;
if projection.offsets.len() < partitions as usize {
projection.offsets.resize(partitions as usize, 0);
}
for (partition, next_offset, messages) in drained {
if projection.offsets[partition as usize] != from[partition as usize] {
continue; }
for message in messages {
let Ok(provenance) = Provenance::try_from(&message) else {
continue;
};
projection.absorb(&message.payload, provenance);
}
projection.offsets[partition as usize] = next_offset;
}
Ok(())
}
pub async fn recall_folded(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
self.catch_up().await?;
let projection = self.projection.lock().await;
let agent_filter = query.agent.as_ref().or(scope.agent.as_ref());
let mut items: Vec<MemoryItem> = projection
.items
.iter()
.filter(|item| {
scope
.conversation
.is_none_or(|conversation| item.provenance.conversation_id == conversation)
})
.filter(|item| !projection.forgotten.contains(&item.id))
.filter(|item| {
agent_filter.is_none_or(|agent| item.provenance.agent.as_ref() == Some(agent))
})
.cloned()
.collect();
if !projection.feedback.is_empty() {
for item in &mut items {
if let Some(weight) = projection.feedback.get(&item.id) {
item.score = Some(*weight);
}
}
items.sort_by(|a, b| {
let a = a.score.unwrap_or(0.0);
let b = b.score.unwrap_or(0.0);
b.total_cmp(&a)
});
items.truncate(query.limit);
} else if items.len() > query.limit {
items = items.split_off(items.len() - query.limit);
}
Ok(items)
}
#[cfg(feature = "kv")]
async fn recall_kv(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
let Some(conversation) = scope.conversation else {
return Ok(Vec::new());
};
let kv = self.laser.kv(self.namespace.clone());
let mut entries = Vec::new();
let mut cursor: Option<Vec<u8>> = None;
loop {
let mut scan = kv.scan().conversation(conversation);
if let Some(resume) = &cursor {
scan = scan.cursor(resume);
}
let page = scan.fetch().await?;
let done = page.cursor.is_none();
cursor = page.cursor;
entries.extend(page.entries);
if done {
break;
}
}
let agent_filter = query.agent.as_ref().or(scope.agent.as_ref());
let mut items: Vec<MemoryItem> = entries
.into_iter()
.filter_map(|entry| self.item_from_entry(conversation, entry))
.filter(|item| {
agent_filter.is_none_or(|agent| item.provenance.agent.as_ref() == Some(agent))
})
.collect();
items.sort_by_key(|item| item.id);
if items.len() > query.limit {
items = items.split_off(items.len() - query.limit);
}
Ok(items)
}
#[cfg(feature = "kv")]
fn item_from_entry(
&self,
conversation: ConversationId,
entry: laser_wire::kv::KvEntry,
) -> Option<MemoryItem> {
let id: MemoryId = entry.key_str()?.parse().ok()?;
let stored = *entry.scope?;
let kind = stored
.kind
.as_deref()
.map(MemoryKind::from_word)
.unwrap_or_default();
let mut provenance = Provenance::builder().conversation_id(conversation).build();
provenance.agent = stored.agent.as_deref().and_then(|agent| agent.parse().ok());
Some(MemoryItem {
id,
payload: entry.value,
provenance,
kind,
source: stored.source,
score: None,
signals: Vec::new(),
})
}
}
impl Memory for LogMemory {
async fn remember(
&self,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.append(scope, MemoryId::new(), MemoryKind::Fact, payload)
.await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
#[cfg(feature = "kv")]
{
self.recall_kv(scope, query).await
}
#[cfg(not(feature = "kv"))]
{
let _ = (scope, query);
Err(LaserError::unsupported(
"memory recall",
"durable recall reads the managed key-value view. Build with the `kv` feature, \
or call `recall_folded` to fold the topic in process",
))
}
}
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
let id = MemoryId::new();
let entry = MemoryLogEntry::Feedback {
target: feedback.target,
weight: feedback.weight,
};
let provenance = Self::provenance(scope, id.to_string());
self.laser
.send_agent_as(
crate::govern::ActionKind::MemoryWrite,
AgentTopic::Custom(&self.topic),
entry.encode()?,
&provenance,
)
.await?;
Ok(id)
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
let entry = MemoryLogEntry::Forget { target: id };
let provenance = Self::provenance(scope, MemoryId::new().to_string());
self.laser
.send_agent_as(
crate::govern::ActionKind::MemoryWrite,
AgentTopic::Custom(&self.topic),
entry.encode()?,
&provenance,
)
.await?;
Ok(())
}
}
#[trait_variant::make(Embedder: Send)]
pub trait LocalEmbedder {
async fn embed(&self, text: &str) -> Result<Vec<f32>, LaserError>;
}
#[trait_variant::make(Reranker: Send)]
pub trait LocalReranker {
async fn rerank(
&self,
query: &str,
items: Vec<MemoryItem>,
) -> Result<Vec<MemoryItem>, LaserError>;
}
#[trait_variant::make(Summarizer: Send)]
pub trait LocalSummarizer {
async fn summarize(&self, bodies: Vec<Vec<u8>>) -> Result<Vec<u8>, LaserError>;
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ConsolidationReport {
pub summarized: usize,
pub reweighted: usize,
pub pruned: usize,
pub derived: usize,
}
#[trait_variant::make(Consolidator: Send)]
pub trait LocalConsolidator {
async fn consolidate(&self, scope: &MemoryScope) -> Result<ConsolidationReport, LaserError>;
}
pub struct DefaultConsolidator<M, S = NoSummarizer> {
memory: M,
max_items: usize,
summarizer: Option<S>,
prune_summarized: bool,
}
pub struct NoSummarizer;
impl Summarizer for NoSummarizer {
async fn summarize(&self, _bodies: Vec<Vec<u8>>) -> Result<Vec<u8>, LaserError> {
Err(LaserError::Config(
"no summarizer was registered on this consolidator",
))
}
}
impl<M> DefaultConsolidator<M> {
pub fn new(memory: M, max_items: usize) -> Self {
Self {
memory,
max_items,
summarizer: None,
prune_summarized: false,
}
}
}
impl<M, S> DefaultConsolidator<M, S> {
pub fn with_summarizer<S2>(self, summarizer: S2) -> DefaultConsolidator<M, S2> {
DefaultConsolidator {
memory: self.memory,
max_items: self.max_items,
summarizer: Some(summarizer),
prune_summarized: self.prune_summarized,
}
}
#[must_use]
pub fn prune_summarized(mut self) -> Self {
self.prune_summarized = true;
self
}
}
impl<M: Memory + Sync, S: Summarizer + Sync> Consolidator for DefaultConsolidator<M, S> {
async fn consolidate(&self, scope: &MemoryScope) -> Result<ConsolidationReport, LaserError> {
let query = MemoryQuery::builder().limit(CONSOLIDATION_WINDOW).build();
let mut items = self.memory.recall(scope, &query).await?;
let mut report = ConsolidationReport::default();
if let Some(summarizer) = &self.summarizer {
let sessions: Vec<&MemoryItem> = items
.iter()
.filter(|item| item.kind == MemoryKind::Message)
.collect();
if !sessions.is_empty() {
let bodies = sessions.iter().map(|item| item.payload.clone()).collect();
let summary = summarizer.summarize(bodies).await?;
self.memory.remember(scope, summary).await?;
report.summarized = sessions.len();
if self.prune_summarized {
let ids: Vec<MemoryId> = sessions.iter().map(|item| item.id).collect();
for id in ids {
if self.memory.forget(scope, id).await.is_ok() {
report.pruned += 1;
}
}
items.retain(|item| item.kind != MemoryKind::Message);
}
}
}
if items.len() <= self.max_items {
return Ok(report);
}
items.sort_by_key(|item| item.id);
let prune_count = items.len() - self.max_items;
for item in items.into_iter().take(prune_count) {
if self.memory.forget(scope, item.id).await.is_ok() {
report.pruned += 1;
}
}
Ok(report)
}
}
pub struct RerankedMemory<M, R> {
inner: M,
reranker: R,
}
impl<M, R> RerankedMemory<M, R> {
pub fn new(inner: M, reranker: R) -> Self {
Self { inner, reranker }
}
}
impl<M: Memory + Sync, R: Reranker + Sync> Memory for RerankedMemory<M, R> {
async fn remember(
&self,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.inner.remember(scope, payload).await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
let candidates = self.inner.recall(scope, query).await?;
match &query.semantic {
Some(text) => self.reranker.rerank(text, candidates).await,
None => Ok(candidates),
}
}
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
self.inner.improve(scope, feedback).await
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
self.inner.forget(scope, id).await
}
}
pub struct VectorMemory<E> {
embedder: E,
items: Mutex<Vec<VectorEntry>>,
governance: Option<VectorGovernance>,
}
#[derive(Clone)]
struct VectorGovernance {
laser: Laser,
}
struct VectorEntry {
id: MemoryId,
scope: MemoryScope,
embedding: Vec<f32>,
item: MemoryItem,
feedback: f32,
}
impl<E> VectorMemory<E> {
pub fn new(embedder: E) -> Self {
Self::with_governance(embedder, None)
}
pub fn governed(laser: Laser, embedder: E) -> Self {
Self::with_governance(embedder, Some(VectorGovernance { laser }))
}
fn with_governance(embedder: E, governance: Option<VectorGovernance>) -> Self {
Self {
embedder,
items: Mutex::new(Vec::new()),
governance,
}
}
async fn govern_entry(
&self,
scope: &MemoryScope,
entry: MemoryLogEntry,
) -> Result<MemoryLogEntry, LaserError> {
let Some(governance) = &self.governance else {
return Ok(entry);
};
let topic = AgentTopic::Audit.topic_string();
let payload = match &entry {
MemoryLogEntry::Item { body, .. } => body.clone(),
MemoryLogEntry::Forget { .. } | MemoryLogEntry::Feedback { .. } => entry.encode()?,
};
let conversation = scope.conversation.unwrap_or_default();
let action = crate::govern::GovernedAction {
kind: crate::govern::ActionKind::MemoryWrite,
stream: governance.laser.stream_required()?,
topic: &topic,
source: scope.agent.as_ref().map(|agent| agent.as_str()),
target: None,
conversation: Some(conversation),
correlation: None,
operation: None,
tool: None,
on_behalf_of: None,
purpose: None,
data_classification: None,
payload: &payload,
signed: false,
counters: crate::govern::ActionCounters::default(),
};
match governance.laser.govern(action).await? {
Some(modified) => match entry {
MemoryLogEntry::Item { id, kind, .. } => Ok(MemoryLogEntry::Item {
id,
kind,
body: modified,
}),
MemoryLogEntry::Forget { .. } | MemoryLogEntry::Feedback { .. } => {
MemoryLogEntry::decode(&modified)
}
},
None => Ok(entry),
}
}
async fn append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError>
where
E: Embedder + Sync,
{
let entry = self
.govern_entry(
scope,
MemoryLogEntry::Item {
id,
kind,
body: payload,
},
)
.await?;
let MemoryLogEntry::Item { id, kind, body } = entry else {
return Err(LaserError::Invalid(
"a memory-item governor modification must remain an item".to_owned(),
));
};
let embedding = self.embedder.embed(&String::from_utf8_lossy(&body)).await?;
let mut provenance = Provenance::builder()
.conversation_id(scope.conversation.unwrap_or_default())
.build();
provenance.agent = scope.agent.clone();
provenance.idempotency_key = Some(id.to_string());
let item = MemoryItem {
id,
payload: body,
provenance,
kind,
score: None,
signals: Vec::new(),
source: None,
};
let mut items = self.items.lock().await;
if items.iter().any(|entry| entry.id == id) {
return Ok(id);
}
items.push(VectorEntry {
id,
scope: scope.clone(),
embedding,
item,
feedback: 0.0,
});
Ok(id)
}
}
impl<E: Embedder + Sync> Memory for VectorMemory<E> {
async fn remember(
&self,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.append(scope, MemoryId::new(), MemoryKind::Fact, payload)
.await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
let wants_semantic = !matches!(query.strategy, RecallStrategy::Keyword);
let wants_keyword = matches!(
query.strategy,
RecallStrategy::Keyword | RecallStrategy::Hybrid
);
let query_embedding = match (&query.semantic, wants_semantic) {
(Some(text), true) => Some(self.embedder.embed(text).await?),
_ => None,
};
let query_tokens = match (&query.semantic, wants_keyword) {
(Some(text), true) => Some(tokenize(text)),
_ => None,
};
let agent_filter = query.agent.as_ref().or(scope.agent.as_ref());
let items = self.items.lock().await;
let matched: Vec<&VectorEntry> = items
.iter()
.filter(|entry| {
scope
.conversation
.is_none_or(|c| entry.scope.conversation == Some(c))
&& agent_filter.is_none_or(|a| entry.scope.agent.as_ref() == Some(a))
&& scope
.user
.as_ref()
.is_none_or(|u| entry.scope.user.as_ref() == Some(u))
&& scope
.app
.as_ref()
.is_none_or(|a| entry.scope.app.as_ref() == Some(a))
})
.collect();
let any_feedback = matched.iter().any(|entry| entry.feedback != 0.0);
let score = |entry: &VectorEntry| {
let mut total = entry.feedback;
if let Some(query_embedding) = &query_embedding {
total += cosine(query_embedding, &entry.embedding);
}
if let Some(query_tokens) = &query_tokens {
total += keyword_score(query_tokens, &entry.item.payload);
}
total
};
let any_query_signal = query_embedding.is_some() || query_tokens.is_some();
if any_query_signal || any_feedback {
let mut scored: Vec<(f32, &VectorEntry)> =
matched.iter().map(|&entry| (score(entry), entry)).collect();
scored.sort_by(|a, b| b.0.total_cmp(&a.0));
Ok(scored
.into_iter()
.take(query.limit)
.map(|(score, entry)| {
let mut item = entry.item.clone();
item.score = Some(score);
item
})
.collect())
} else {
let start = matched.len().saturating_sub(query.limit);
Ok(matched[start..]
.iter()
.map(|entry| entry.item.clone())
.collect())
}
}
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
let entry = self
.govern_entry(
scope,
MemoryLogEntry::Feedback {
target: feedback.target,
weight: feedback.weight,
},
)
.await?;
let MemoryLogEntry::Feedback { target, weight } = entry else {
return Err(LaserError::Invalid(
"a memory-feedback governor modification must remain feedback".to_owned(),
));
};
let mut items = self.items.lock().await;
if let Some(entry) = items.iter_mut().find(|entry| entry.id == target) {
entry.feedback += weight;
}
Ok(MemoryId::new())
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
let entry = self
.govern_entry(scope, MemoryLogEntry::Forget { target: id })
.await?;
let MemoryLogEntry::Forget { target } = entry else {
return Err(LaserError::Invalid(
"a memory-forget governor modification must remain a tombstone".to_owned(),
));
};
self.items.lock().await.retain(|entry| entry.id != target);
Ok(())
}
}
fn cosine(a: &[f32], b: &[f32]) -> f32 {
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let norm_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let norm_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
0.0
} else {
dot / (norm_a * norm_b)
}
}
fn tokenize(text: &str) -> HashSet<String> {
text.split(|c: char| !c.is_alphanumeric())
.filter(|token| !token.is_empty())
.map(str::to_ascii_lowercase)
.collect()
}
fn put_header(
headers: &mut std::collections::BTreeMap<HeaderKey, HeaderValue>,
key: &str,
value: &str,
) -> Result<(), LaserError> {
let key = HeaderKey::from_str(key)
.map_err(|error| LaserError::Codec(format!("memory header key {key}: {error}")))?;
let value = HeaderValue::from_str(value)
.map_err(|error| LaserError::Codec(format!("memory header value: {error}")))?;
headers.insert(key, value);
Ok(())
}
pub fn fuse_reciprocal_rank(signals: Vec<Vec<MemoryItem>>, limit: usize) -> Vec<MemoryItem> {
const RRF_K: f32 = 60.0;
let mut fused: Vec<MemoryItem> = Vec::new();
for ranked in signals {
for item in ranked {
let contribution = item
.signals
.first()
.map_or(0.0, |signal| 1.0 / (RRF_K + signal.rank as f32));
match fused.iter_mut().find(|held| held.id == item.id) {
Some(held) => {
held.score = Some(held.score.unwrap_or(0.0) + contribution);
held.signals.extend(item.signals);
}
None => {
let mut item = item;
item.score = Some(contribution);
fused.push(item);
}
}
}
}
fused.sort_by(|a, b| b.score.unwrap_or(0.0).total_cmp(&a.score.unwrap_or(0.0)));
fused.truncate(limit);
fused
}
fn keyword_score(query_tokens: &HashSet<String>, body: &[u8]) -> f32 {
if query_tokens.is_empty() {
return 0.0;
}
let Ok(text) = std::str::from_utf8(body) else {
return 0.0;
};
let body_tokens = tokenize(text);
let hits = query_tokens
.iter()
.filter(|token| body_tokens.contains(*token))
.count();
hits as f32 / query_tokens.len() as f32
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum MemoryBackend {
#[default]
Auto,
Log,
Vector,
}
pub struct MemoryTopicBuilder {
laser: Laser,
stream: Option<String>,
topic: String,
partitions: u32,
ttl: Option<std::time::Duration>,
}
impl MemoryTopicBuilder {
#[must_use]
pub fn stream(mut self, stream: impl Into<String>) -> Self {
self.stream = Some(stream.into());
self
}
#[must_use]
pub fn partitions(mut self, partitions: u32) -> Self {
self.partitions = partitions.max(1);
self
}
#[must_use]
pub fn ttl(mut self, ttl: std::time::Duration) -> Self {
self.ttl = Some(ttl);
self
}
#[must_use]
pub fn no_expiry(mut self) -> Self {
self.ttl = None;
self
}
pub async fn build(self) -> Result<MemoryHandle, LaserError> {
let stream = match &self.stream {
Some(stream) => stream.clone(),
None => self.laser.stream_required()?.to_owned(),
};
let expiry = match self.ttl {
Some(ttl) => IggyExpiry::ExpireDuration(IggyDuration::from(ttl)),
None => IggyExpiry::NeverExpire,
};
self.laser
.ensure_topic_on_with(&stream, &self.topic, self.partitions, expiry)
.await?;
let topic = Identifier::named(&self.topic)?;
Ok(MemoryHandle::Log(LogMemory::on_stream_topic(
self.laser,
self.stream,
topic,
)))
}
}
#[allow(clippy::large_enum_variant)]
pub enum MemoryHandle {
Log(LogMemory),
Vector(std::sync::Arc<VectorMemory<SharedEmbedder>>),
Reranked {
inner: Box<MemoryHandle>,
reranker: SharedReranker,
},
Custom(std::sync::Arc<dyn DynMemory>),
}
#[derive(Clone)]
pub struct SharedEmbedder(std::sync::Arc<dyn DynEmbedder>);
impl Embedder for SharedEmbedder {
async fn embed(&self, text: &str) -> Result<Vec<f32>, LaserError> {
self.0.embed_dyn(text).await
}
}
trait DynEmbedder: Send + Sync {
fn embed_dyn<'a>(
&'a self,
text: &'a str,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<f32>, LaserError>> + Send + 'a>>;
}
impl<E: Embedder + Sync> DynEmbedder for E {
fn embed_dyn<'a>(
&'a self,
text: &'a str,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<f32>, LaserError>> + Send + 'a>> {
Box::pin(self.embed(text))
}
}
#[derive(Clone)]
pub struct SharedReranker(std::sync::Arc<dyn DynReranker>);
impl Reranker for SharedReranker {
async fn rerank(
&self,
query: &str,
items: Vec<MemoryItem>,
) -> Result<Vec<MemoryItem>, LaserError> {
self.0.rerank_dyn(query, items).await
}
}
trait DynReranker: Send + Sync {
fn rerank_dyn<'a>(
&'a self,
query: &'a str,
items: Vec<MemoryItem>,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>>;
}
impl<R: Reranker + Sync> DynReranker for R {
fn rerank_dyn<'a>(
&'a self,
query: &'a str,
items: Vec<MemoryItem>,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>>
{
Box::pin(self.rerank(query, items))
}
}
impl MemoryHandle {
#[must_use]
pub fn backend(&self) -> &'static str {
match self {
Self::Log(_) => "log",
Self::Vector(_) => "vector",
Self::Reranked { inner, .. } => inner.backend(),
Self::Custom(_) => "custom",
}
}
#[must_use]
pub fn embedder(self, embedder: impl Embedder + Sync + 'static) -> Self {
self.embedder_shared(SharedEmbedder(std::sync::Arc::new(embedder)))
}
fn embedder_shared(self, shared: SharedEmbedder) -> Self {
match self {
Self::Vector(memory) => Self::Vector(std::sync::Arc::new(
VectorMemory::with_governance(shared, memory.governance.clone()),
)),
Self::Reranked { inner, reranker } => Self::Reranked {
inner: Box::new(inner.embedder_shared(shared)),
reranker,
},
Self::Log(memory) => Self::Log(memory),
Self::Custom(memory) => Self::Custom(memory),
}
}
#[must_use]
pub fn reranker(self, reranker: impl Reranker + Sync + 'static) -> Self {
Self::Reranked {
inner: Box::new(self),
reranker: SharedReranker(std::sync::Arc::new(reranker)),
}
}
}
impl Laser {
pub fn memory(&self, namespace: impl Into<String>) -> MemoryHandle {
self.memory_with(namespace, MemoryBackend::Auto)
}
pub fn memory_on_topic(&self, topic: impl AsRef<str>) -> Result<MemoryHandle, LaserError> {
let topic = Identifier::named(topic.as_ref())?;
Ok(MemoryHandle::Log(LogMemory::on_topic(self.clone(), topic)))
}
pub fn memory_topic(&self, topic: impl Into<String>) -> MemoryTopicBuilder {
MemoryTopicBuilder {
laser: self.clone(),
stream: None,
topic: topic.into(),
partitions: 1,
ttl: Some(DEFAULT_MEMORY_TOPIC_TTL),
}
}
pub fn memory_with(
&self,
namespace: impl Into<String>,
backend: MemoryBackend,
) -> MemoryHandle {
let namespace = namespace.into();
match backend {
MemoryBackend::Auto | MemoryBackend::Log => {
MemoryHandle::Log(LogMemory::in_namespace(self.clone(), namespace))
}
MemoryBackend::Vector => {
let _ = namespace;
MemoryHandle::Vector(std::sync::Arc::new(VectorMemory::governed(
self.clone(),
SharedEmbedder(std::sync::Arc::new(MaybeEmbedder(None))),
)))
}
}
}
pub fn memory_custom(
&self,
backend: std::sync::Arc<dyn crate::memory::DynMemory>,
) -> MemoryHandle {
MemoryHandle::Custom(backend)
}
}
impl MemoryHandle {
async fn append_id(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
match self {
MemoryHandle::Log(memory) => memory.append(scope, id, kind, payload).await,
MemoryHandle::Vector(memory) => memory.append(scope, id, kind, payload).await,
MemoryHandle::Reranked { inner, .. } => {
append_id_boxed(inner, scope, id, kind, payload).await
}
MemoryHandle::Custom(memory) => memory.remember(scope, payload).await,
}
}
pub async fn set(
&self,
key: impl AsRef<str>,
payload: impl Into<Vec<u8>>,
) -> Result<(), LaserError> {
self.named_memory("set")?
.set_named(key.as_ref(), payload.into())
.await
}
pub async fn fetch(&self, key: impl AsRef<str>) -> Result<Option<Vec<u8>>, LaserError> {
self.named_memory("fetch")?.fetch_named(key.as_ref()).await
}
pub async fn fetch_folded(&self, key: impl AsRef<str>) -> Result<Option<Vec<u8>>, LaserError> {
self.named_memory("fetch_folded")?
.fetch_named_folded(key.as_ref())
.await
}
pub async fn update(
&self,
key: impl AsRef<str>,
patch: impl Into<Vec<u8>>,
) -> Result<(), LaserError> {
self.named_memory("update")?
.update_named(key.as_ref(), patch.into())
.await
}
pub async fn remove(&self, key: impl AsRef<str>) -> Result<(), LaserError> {
self.named_memory("remove")?
.forget_named(key.as_ref())
.await
}
fn named_memory(&self, verb: &str) -> Result<&LogMemory, LaserError> {
match self {
MemoryHandle::Log(memory) => Ok(memory),
MemoryHandle::Reranked { inner, .. } => inner.named_memory(verb),
other => Err(LaserError::unsupported(
"memory",
format!(
"{verb}(key) is the named-item altitude and needs the durable memory handle \
(this handle resolved to `{}`)",
other.backend()
),
)),
}
}
pub fn remember(&self, payload: impl Into<Vec<u8>>) -> RememberBuilder<'_> {
RememberBuilder {
handle: self,
payload: payload.into(),
scope: MemoryScope::default(),
kind: MemoryKind::Fact,
dedup: false,
}
}
pub async fn recall_folded(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
match self {
MemoryHandle::Log(memory) => memory.recall_folded(scope, query).await,
MemoryHandle::Vector(memory) => Memory::recall(memory.as_ref(), scope, query).await,
MemoryHandle::Reranked { inner, .. } => recall_folded_boxed(inner, scope, query).await,
MemoryHandle::Custom(memory) => memory.recall(scope, query).await,
}
}
pub fn recall(&self, conversation: ConversationId) -> RecallBuilder<'_> {
RecallBuilder {
handle: self,
scope: MemoryScope::builder().conversation(conversation).build(),
limit: 50,
semantic: None,
strategy: RecallStrategy::Auto,
folded: false,
}
}
pub async fn context(
&self,
conversation: ConversationId,
token_budget: Option<usize>,
) -> Result<String, LaserError> {
let items = self.recall(conversation).fetch().await?;
Ok(to_context_block(&items, token_budget))
}
pub async fn consolidate(
&self,
scope: &MemoryScope,
max_items: usize,
) -> Result<ConsolidationReport, LaserError> {
Consolidator::consolidate(&DefaultConsolidator::new(self, max_items), scope).await
}
pub async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
Memory::improve(self, scope, feedback).await
}
pub async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
Memory::forget(self, scope, id).await
}
}
impl Memory for MemoryHandle {
async fn remember(
&self,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
if let MemoryHandle::Custom(memory) = self {
return memory.remember(scope, payload).await;
}
self.append_id(scope, MemoryId::new(), MemoryKind::Fact, payload)
.await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
match self {
MemoryHandle::Log(memory) => Memory::recall(memory, scope, query).await,
MemoryHandle::Vector(memory) => Memory::recall(memory.as_ref(), scope, query).await,
MemoryHandle::Reranked { inner, reranker } => {
let items = recall_boxed(inner, scope, query).await?;
match &query.semantic {
Some(text) => Reranker::rerank(reranker, text, items).await,
None => Ok(items),
}
}
MemoryHandle::Custom(memory) => memory.recall(scope, query).await,
}
}
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
match self {
MemoryHandle::Log(memory) => Memory::improve(memory, scope, feedback).await,
MemoryHandle::Vector(memory) => Memory::improve(memory.as_ref(), scope, feedback).await,
MemoryHandle::Reranked { inner, .. } => improve_boxed(inner, scope, feedback).await,
MemoryHandle::Custom(memory) => memory.improve(scope, feedback).await,
}
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
match self {
MemoryHandle::Log(memory) => Memory::forget(memory, scope, id).await,
MemoryHandle::Vector(memory) => Memory::forget(memory.as_ref(), scope, id).await,
MemoryHandle::Reranked { inner, .. } => forget_boxed(inner, scope, id).await,
MemoryHandle::Custom(memory) => memory.forget(scope, id).await,
}
}
}
impl Memory for &MemoryHandle {
async fn remember(
&self,
scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
Memory::remember(*self, scope, payload).await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
Memory::recall(*self, scope, query).await
}
async fn improve(
&self,
scope: &MemoryScope,
feedback: Feedback,
) -> Result<MemoryId, LaserError> {
Memory::improve(*self, scope, feedback).await
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
Memory::forget(*self, scope, id).await
}
}
#[derive(Clone)]
pub struct SharedConsolidator(std::sync::Arc<dyn DynConsolidator>);
impl SharedConsolidator {
pub fn new(consolidator: impl Consolidator + Sync + 'static) -> Self {
Self(std::sync::Arc::new(consolidator))
}
}
impl Consolidator for SharedConsolidator {
async fn consolidate(&self, scope: &MemoryScope) -> Result<ConsolidationReport, LaserError> {
self.0.consolidate_dyn(scope).await
}
}
trait DynConsolidator: Send + Sync {
fn consolidate_dyn<'a>(
&'a self,
scope: &'a MemoryScope,
) -> std::pin::Pin<Box<dyn Future<Output = Result<ConsolidationReport, LaserError>> + Send + 'a>>;
}
impl<C: Consolidator + Sync> DynConsolidator for C {
fn consolidate_dyn<'a>(
&'a self,
scope: &'a MemoryScope,
) -> std::pin::Pin<Box<dyn Future<Output = Result<ConsolidationReport, LaserError>> + Send + 'a>>
{
Box::pin(Consolidator::consolidate(self, scope))
}
}
fn append_id_boxed<'a>(
inner: &'a MemoryHandle,
scope: &'a MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> std::pin::Pin<Box<dyn Future<Output = Result<MemoryId, LaserError>> + Send + 'a>> {
Box::pin(inner.append_id(scope, id, kind, payload))
}
fn recall_boxed<'a>(
inner: &'a MemoryHandle,
scope: &'a MemoryScope,
query: &'a MemoryQuery,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>> {
Box::pin(Memory::recall(inner, scope, query))
}
fn recall_folded_boxed<'a>(
inner: &'a MemoryHandle,
scope: &'a MemoryScope,
query: &'a MemoryQuery,
) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, LaserError>> + Send + 'a>> {
Box::pin(inner.recall_folded(scope, query))
}
fn improve_boxed<'a>(
inner: &'a MemoryHandle,
scope: &'a MemoryScope,
feedback: Feedback,
) -> std::pin::Pin<Box<dyn Future<Output = Result<MemoryId, LaserError>> + Send + 'a>> {
Box::pin(Memory::improve(inner, scope, feedback))
}
fn forget_boxed<'a>(
inner: &'a MemoryHandle,
scope: &'a MemoryScope,
id: MemoryId,
) -> std::pin::Pin<Box<dyn Future<Output = Result<(), LaserError>> + Send + 'a>> {
Box::pin(Memory::forget(inner, scope, id))
}
fn merge_named(current: Option<&[u8]>, patch: &[u8]) -> Result<Vec<u8>, LaserError> {
let patch: serde_json::Value = serde_json::from_slice(patch)
.map_err(|error| LaserError::Codec(format!("update patch is not JSON: {error}")))?;
let mut base: serde_json::Value = match current {
Some(payload) => serde_json::from_slice(payload)
.map_err(|error| LaserError::Codec(format!("stored value is not JSON: {error}")))?,
None => serde_json::Value::Null,
};
merge_patch(&mut base, &patch);
serde_json::to_vec(&base)
.map_err(|error| LaserError::Codec(format!("encode merged value: {error}")))
}
fn merge_patch(base: &mut serde_json::Value, patch: &serde_json::Value) {
let serde_json::Value::Object(fields) = patch else {
*base = patch.clone();
return;
};
if !base.is_object() {
*base = serde_json::Value::Object(serde_json::Map::new());
}
let map = base.as_object_mut().expect("base was just made an object");
for (key, value) in fields {
if value.is_null() {
map.remove(key);
} else {
merge_patch(
map.entry(key.clone()).or_insert(serde_json::Value::Null),
value,
);
}
}
}
#[derive(Clone)]
pub struct MaybeEmbedder(Option<SharedEmbedder>);
impl Embedder for MaybeEmbedder {
async fn embed(&self, text: &str) -> Result<Vec<f32>, LaserError> {
match &self.0 {
Some(embedder) => Embedder::embed(embedder, text).await,
None => Err(LaserError::Config(
"this memory needs an embedder for semantic recall: register one with \
MemoryHandle::embedder(..)",
)),
}
}
}
pub struct RememberBuilder<'a> {
handle: &'a MemoryHandle,
payload: Vec<u8>,
scope: MemoryScope,
kind: MemoryKind,
dedup: bool,
}
impl RememberBuilder<'_> {
#[must_use]
pub fn scope(mut self, conversation: ConversationId) -> Self {
self.scope.conversation = Some(conversation);
self
}
#[must_use]
pub fn agent(mut self, agent: AgentId) -> Self {
self.scope.agent = Some(agent);
self
}
#[must_use]
pub fn stream(mut self, stream: impl Into<String>) -> Self {
self.scope.stream = Some(stream.into());
self
}
#[must_use]
pub fn durable(mut self) -> Self {
self.scope.lifetime = Lifetime::Durable;
self
}
#[must_use]
pub fn kind(mut self, kind: MemoryKind) -> Self {
self.kind = kind;
self
}
#[must_use]
pub fn dedup(mut self) -> Self {
self.dedup = true;
self
}
pub async fn send(self) -> Result<MemoryId, LaserError> {
let id = if self.dedup {
MemoryId::content(&self.scope, self.kind, &self.payload)
} else {
MemoryId::new()
};
self.handle
.append_id(&self.scope, id, self.kind, self.payload)
.await
}
}
pub struct RecallBuilder<'a> {
handle: &'a MemoryHandle,
scope: MemoryScope,
limit: usize,
semantic: Option<String>,
strategy: RecallStrategy,
folded: bool,
}
impl RecallBuilder<'_> {
#[must_use]
pub fn semantic(mut self, text: impl Into<String>) -> Self {
self.semantic = Some(text.into());
self.strategy = RecallStrategy::Semantic;
self
}
#[must_use]
pub fn recent(mut self) -> Self {
self.strategy = RecallStrategy::Recent;
self
}
#[must_use]
pub fn folded(mut self) -> Self {
self.folded = true;
self
}
#[must_use]
pub fn keyword(mut self, text: impl Into<String>) -> Self {
self.semantic = Some(text.into());
self.strategy = RecallStrategy::Keyword;
self
}
#[must_use]
pub fn hybrid(mut self, text: impl Into<String>) -> Self {
self.semantic = Some(text.into());
self.strategy = RecallStrategy::Hybrid;
self
}
#[must_use]
pub fn strategy(mut self, strategy: RecallStrategy) -> Self {
self.strategy = strategy;
self
}
#[must_use]
pub fn agent(mut self, agent: AgentId) -> Self {
self.scope.agent = Some(agent);
self
}
#[must_use]
pub fn limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
pub async fn fetch(self) -> Result<Vec<MemoryItem>, LaserError> {
let query = MemoryQuery::builder()
.limit(self.limit)
.maybe_semantic(self.semantic)
.strategy(self.strategy)
.maybe_agent(self.scope.agent.clone())
.build();
if self.folded {
self.handle.recall_folded(&self.scope, &query).await
} else {
Memory::recall(self.handle, &self.scope, &query).await
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn given_a_custom_memory_backend_when_used_through_the_handle_then_should_delegate() {
struct FakeMemory(std::sync::Mutex<Vec<Vec<u8>>>);
impl Memory for FakeMemory {
async fn remember(
&self,
_scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.0.lock().expect("lock").push(payload);
Ok(MemoryId::new())
}
async fn recall(
&self,
_scope: &MemoryScope,
_query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
let provenance = Provenance::builder()
.conversation_id(ConversationId::new())
.build();
Ok(self
.0
.lock()
.expect("lock")
.iter()
.map(|payload| {
MemoryItem::plain(MemoryId::new(), payload.clone(), provenance.clone())
})
.collect())
}
async fn improve(
&self,
_scope: &MemoryScope,
_feedback: Feedback,
) -> Result<MemoryId, LaserError> {
Ok(MemoryId::new())
}
async fn forget(&self, _scope: &MemoryScope, _id: MemoryId) -> Result<(), LaserError> {
Ok(())
}
}
let handle = MemoryHandle::Custom(std::sync::Arc::new(FakeMemory(std::sync::Mutex::new(
Vec::new(),
))));
let scope = MemoryScope::default();
Memory::remember(&handle, &scope, b"hi".to_vec())
.await
.expect("custom remember");
let items = Memory::recall(&handle, &scope, &MemoryQuery::builder().limit(10).build())
.await
.expect("custom recall");
assert_eq!(items.len(), 1);
assert_eq!(items[0].payload, b"hi");
}
#[test]
fn given_a_log_entry_when_encoded_then_the_wire_memory_record_decodes_it() {
let id = MemoryId::new();
let entry = MemoryLogEntry::Item {
id,
kind: MemoryKind::Message,
body: b"checkout is slow".to_vec(),
};
let payload = entry.encode().expect("entry encodes");
let record: laser_wire::memory::MemoryRecord =
laser_wire::framing::decode_named(&payload).expect("wire record decodes the SDK entry");
match record {
laser_wire::memory::MemoryRecord::Item {
id: rid,
kind,
body,
} => {
assert_eq!(rid, id.to_string());
assert_eq!(kind, "message");
assert_eq!(body, b"checkout is slow");
}
other => panic!("expected an Item record, got {other:?}"),
}
}
#[test]
fn given_a_memory_id_when_round_tripped_through_a_string_then_should_be_equal() {
let id = MemoryId::new();
let parsed = id
.to_string()
.parse::<MemoryId>()
.expect("a formatted memory id should parse");
assert_eq!(parsed, id);
}
struct MockMemory {
items: std::sync::Mutex<Vec<MemoryItem>>,
}
impl Memory for MockMemory {
async fn remember(
&self,
_scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
let id = MemoryId::new();
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
self.items
.lock()
.expect("lock")
.push(MemoryItem::plain(id, payload, provenance));
Ok(id)
}
async fn recall(
&self,
_scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
Ok(self
.items
.lock()
.expect("lock")
.iter()
.take(query.limit)
.cloned()
.collect())
}
async fn improve(
&self,
_scope: &MemoryScope,
_feedback: Feedback,
) -> Result<MemoryId, LaserError> {
Ok(MemoryId::new())
}
async fn forget(&self, _scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
self.items
.lock()
.expect("lock")
.retain(|item| item.id != id);
Ok(())
}
}
#[test]
fn given_more_than_max_items_when_consolidated_then_should_prune_down_to_the_cap() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.expect("runtime builds");
runtime.block_on(async {
let mock = MockMemory {
items: std::sync::Mutex::new(Vec::new()),
};
let scope = MemoryScope::builder().build();
for i in 0..5 {
Memory::remember(&mock, &scope, format!("item-{i}").into_bytes())
.await
.expect("remember");
}
let consolidator = DefaultConsolidator::new(mock, 2);
let report = Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidate");
assert_eq!(
report.pruned, 3,
"five items down to a cap of two prunes three"
);
let remaining = Memory::recall(
&consolidator.memory,
&scope,
&MemoryQuery::builder().limit(100).build(),
)
.await
.expect("recall");
assert_eq!(remaining.len(), 2);
});
}
struct StubSummarizer;
impl Summarizer for StubSummarizer {
async fn summarize(&self, bodies: Vec<Vec<u8>>) -> Result<Vec<u8>, LaserError> {
Ok(format!("summary of {}", bodies.len()).into_bytes())
}
}
#[tokio::test]
async fn given_session_messages_when_consolidated_then_should_store_one_summary() {
let memory = LogMemoryStub::default();
let scope = MemoryScope::default();
for body in [b"turn one".to_vec(), b"turn two".to_vec()] {
memory.push(MemoryKind::Message, body);
}
memory.push(MemoryKind::Fact, b"a durable fact".to_vec());
let consolidator =
DefaultConsolidator::new(memory.clone(), 100).with_summarizer(StubSummarizer);
let report = Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidates");
assert_eq!(report.summarized, 2, "both session turns fold");
assert_eq!(report.pruned, 0, "prune_summarized was not opted in");
assert!(
memory.bodies().iter().any(|body| body == b"summary of 2"),
"the summary landed"
);
}
#[tokio::test]
async fn given_prune_summarized_when_consolidated_then_should_forget_the_turns() {
let memory = LogMemoryStub::default();
let scope = MemoryScope::default();
memory.push(MemoryKind::Message, b"turn".to_vec());
let consolidator = DefaultConsolidator::new(memory.clone(), 100)
.with_summarizer(StubSummarizer)
.prune_summarized();
let report = Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidates");
assert_eq!(report.summarized, 1);
assert_eq!(report.pruned, 1, "the folded turn is forgotten");
}
#[derive(Clone, Default)]
struct LogMemoryStub {
items: std::sync::Arc<std::sync::Mutex<Vec<MemoryItem>>>,
}
impl LogMemoryStub {
fn push(&self, kind: MemoryKind, payload: Vec<u8>) {
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
let mut item = MemoryItem::plain(MemoryId::new(), payload, provenance);
item.kind = kind;
self.items.lock().expect("stub lock").push(item);
}
fn bodies(&self) -> Vec<Vec<u8>> {
self.items
.lock()
.expect("stub lock")
.iter()
.map(|item| item.payload.clone())
.collect()
}
}
impl Memory for LogMemoryStub {
async fn remember(
&self,
_scope: &MemoryScope,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
let id = MemoryId::new();
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
self.items
.lock()
.expect("stub lock")
.push(MemoryItem::plain(id, payload, provenance));
Ok(id)
}
async fn recall(
&self,
_scope: &MemoryScope,
_query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
Ok(self.items.lock().expect("stub lock").clone())
}
async fn improve(
&self,
_scope: &MemoryScope,
_feedback: Feedback,
) -> Result<MemoryId, LaserError> {
Ok(MemoryId::new())
}
async fn forget(&self, _scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
self.items
.lock()
.expect("stub lock")
.retain(|item| item.id != id);
Ok(())
}
}
fn ranked_item(id: MemoryId, strategy: RecallStrategy, rank: usize) -> MemoryItem {
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
let mut item = MemoryItem::plain(id, b"body".to_vec(), provenance);
item.signals = vec![RecallSignal {
strategy,
rank,
score: None,
}];
item
}
#[test]
fn given_two_signals_when_fused_then_should_rank_agreement_first() {
let shared = MemoryId::new();
let semantic_only = MemoryId::new();
let keyword_only = MemoryId::new();
let semantic = vec![
ranked_item(semantic_only, RecallStrategy::Semantic, 0),
ranked_item(shared, RecallStrategy::Semantic, 1),
];
let keyword = vec![
ranked_item(keyword_only, RecallStrategy::Keyword, 0),
ranked_item(shared, RecallStrategy::Keyword, 1),
];
let fused = fuse_reciprocal_rank(vec![semantic, keyword], 10);
assert_eq!(fused[0].id, shared, "both signals agree on it");
assert_eq!(fused[0].signals.len(), 2, "attribution keeps both signals");
assert_eq!(fused.len(), 3);
}
#[test]
fn given_a_limit_when_fused_then_should_truncate_after_ranking() {
let lists = vec![vec![
ranked_item(MemoryId::new(), RecallStrategy::Semantic, 0),
ranked_item(MemoryId::new(), RecallStrategy::Semantic, 1),
ranked_item(MemoryId::new(), RecallStrategy::Semantic, 2),
]];
assert_eq!(fuse_reciprocal_rank(lists, 2).len(), 2);
}
fn context_item(body: &str) -> MemoryItem {
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
MemoryItem::plain(MemoryId::new(), body.as_bytes().to_vec(), provenance)
}
#[test]
fn given_no_budget_when_rendered_then_should_join_every_item() {
let items = [context_item("alpha"), context_item("beta")];
assert_eq!(to_context_block(&items, None), "alpha\n\nbeta");
}
#[test]
fn given_a_token_budget_when_rendered_then_should_fill_in_order_and_mark_the_overflow() {
let items = [
context_item("first twenty bytes!!"),
context_item("second twenty bytes!"),
];
let block = to_context_block(&items, Some(6));
assert_eq!(
block,
"first twenty bytes!!\n\n[... 1 more recalled item(s) omitted ...]"
);
}
#[test]
fn given_a_budget_smaller_than_the_first_item_when_rendered_then_should_keep_at_least_one() {
let items = [context_item("a much longer single memory item")];
let block = to_context_block(&items, Some(1));
assert_eq!(block, "a much longer single memory item");
}
fn item_entry(id: MemoryId, body: &[u8]) -> Vec<u8> {
MemoryLogEntry::Item {
id,
kind: MemoryKind::Fact,
body: body.to_vec(),
}
.encode()
.expect("an item entry encodes")
}
#[test]
fn given_messages_when_absorbed_then_should_fold_items_and_apply_tombstones() {
let conversation = ConversationId::new();
let prov = || Provenance::builder().conversation_id(conversation).build();
let mut projection = Projection::default();
let id_a = MemoryId::new();
let id_b = MemoryId::new();
projection.absorb(&item_entry(id_a, b"a"), prov());
projection.absorb(&item_entry(id_b, b"b"), prov());
assert_eq!(projection.items.len(), 2, "two items folded in");
let tombstone = MemoryLogEntry::Forget { target: id_a }
.encode()
.expect("a forget entry encodes");
projection.absorb(&tombstone, prov());
assert_eq!(projection.items.len(), 1, "tombstoned item reclaimed");
assert_eq!(projection.items[0].id, id_b, "only the survivor remains");
assert!(projection.forgotten.contains(&id_a), "first item forgotten");
projection.absorb(&item_entry(id_a, b"a-again"), prov());
assert_eq!(projection.items.len(), 1, "forgotten id not re-added");
projection.absorb(b"not an entry", prov());
assert_eq!(projection.items.len(), 1, "non-entry payload dropped");
}
#[test]
fn given_a_feedback_entry_when_absorbed_then_should_accumulate_the_weight() {
let conversation = ConversationId::new();
let prov = || Provenance::builder().conversation_id(conversation).build();
let mut projection = Projection::default();
let target = MemoryId::new();
let feedback = MemoryLogEntry::Feedback {
target,
weight: 2.5,
}
.encode()
.expect("a feedback entry encodes");
projection.absorb(&feedback, prov());
projection.absorb(&feedback, prov());
assert_eq!(projection.feedback.get(&target), Some(&5.0));
assert!(projection.items.is_empty(), "feedback is not an item");
}
struct WordEmbedder;
impl Embedder for WordEmbedder {
async fn embed(&self, text: &str) -> Result<Vec<f32>, LaserError> {
Ok(vec![
text.matches("cat").count() as f32,
text.matches("dog").count() as f32,
text.matches("fish").count() as f32,
])
}
}
#[tokio::test]
async fn given_a_semantic_query_when_recalling_then_should_rank_by_similarity() {
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::builder()
.conversation(ConversationId::new())
.build();
Memory::remember(&memory, &scope, b"the cat sat".to_vec())
.await
.expect("remembering the cat note should succeed");
Memory::remember(&memory, &scope, b"the dog ran".to_vec())
.await
.expect("remembering the dog note should succeed");
let top = Memory::recall(
&memory,
&scope,
&MemoryQuery::builder()
.semantic("cat".to_owned())
.limit(1)
.build(),
)
.await
.expect("recall should succeed");
assert_eq!(top.len(), 1);
assert_eq!(top[0].payload.as_slice(), b"the cat sat");
}
#[tokio::test]
async fn given_a_deduped_vector_remember_when_repeated_then_should_store_one_item() {
let handle = MemoryHandle::Vector(std::sync::Arc::new(VectorMemory::new(SharedEmbedder(
std::sync::Arc::new(WordEmbedder),
))));
let conversation = ConversationId::new();
let first = handle
.remember(b"stable fact".to_vec())
.scope(conversation)
.dedup()
.send()
.await
.expect("first remember succeeds");
let second = handle
.remember(b"stable fact".to_vec())
.scope(conversation)
.dedup()
.send()
.await
.expect("second remember succeeds");
let recalled = handle
.recall(conversation)
.fetch()
.await
.expect("vector recall succeeds");
assert_eq!(first, second);
assert_eq!(recalled.len(), 1);
}
#[test]
fn given_a_memory_backend_when_defaulted_then_should_be_auto() {
assert_eq!(MemoryBackend::default(), MemoryBackend::Auto);
}
#[test]
fn given_memory_kinds_when_classified_then_should_map_to_the_taxonomy() {
assert_eq!(MemoryKind::Message.class(), MemoryClass::Episodic);
assert_eq!(MemoryKind::Procedure.class(), MemoryClass::Procedural);
assert_eq!(MemoryKind::Fact.class(), MemoryClass::Semantic);
assert_eq!(MemoryKind::Summary.class(), MemoryClass::Semantic);
}
#[tokio::test]
async fn given_a_keyword_strategy_when_recalling_then_should_match_exact_terms() {
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::builder()
.conversation(ConversationId::new())
.build();
Memory::remember(&memory, &scope, b"alpha beta".to_vec())
.await
.expect("remember alpha");
Memory::remember(&memory, &scope, b"gamma delta".to_vec())
.await
.expect("remember gamma");
let top = Memory::recall(
&memory,
&scope,
&MemoryQuery::builder()
.semantic("gamma".to_owned())
.strategy(RecallStrategy::Keyword)
.limit(1)
.build(),
)
.await
.expect("keyword recall");
assert_eq!(top.len(), 1);
assert_eq!(top[0].payload.as_slice(), b"gamma delta");
}
#[tokio::test]
async fn given_a_user_scope_when_recalling_then_should_narrow_to_that_user() {
let memory = VectorMemory::new(WordEmbedder);
let conversation = ConversationId::new();
let alice = MemoryScope::builder()
.conversation(conversation)
.user("alice".to_owned())
.build();
let bob = MemoryScope::builder()
.conversation(conversation)
.user("bob".to_owned())
.build();
Memory::remember(&memory, &alice, b"alice secret".to_vec())
.await
.expect("remember alice");
Memory::remember(&memory, &bob, b"bob secret".to_vec())
.await
.expect("remember bob");
let hits = Memory::recall(&memory, &alice, &MemoryQuery::builder().build())
.await
.expect("recall alice");
assert_eq!(hits.len(), 1, "only alice's item is in scope");
assert_eq!(hits[0].payload.as_slice(), b"alice secret");
}
struct ReverseReranker;
impl Reranker for ReverseReranker {
async fn rerank(
&self,
_query: &str,
mut items: Vec<MemoryItem>,
) -> Result<Vec<MemoryItem>, LaserError> {
items.reverse();
Ok(items)
}
}
#[tokio::test]
async fn given_a_reranked_memory_when_recalling_then_should_apply_the_reranker() {
let inner = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::builder()
.conversation(ConversationId::new())
.build();
Memory::remember(&inner, &scope, b"the cat sat".to_vec())
.await
.expect("remember cat");
Memory::remember(&inner, &scope, b"the dog ran".to_vec())
.await
.expect("remember dog");
let reranked = RerankedMemory::new(inner, ReverseReranker);
let hits = Memory::recall(
&reranked,
&scope,
&MemoryQuery::builder().semantic("cat".to_owned()).build(),
)
.await
.expect("reranked recall");
assert_eq!(hits.len(), 2);
assert_eq!(
hits[0].payload.as_slice(),
b"the dog ran",
"the reranker moved the dog note to the front"
);
}
#[test]
fn given_the_same_content_when_addressed_then_should_produce_the_same_id() {
let owner = MemoryScope::builder()
.agent("planner".parse().expect("valid agent"))
.lifetime(Lifetime::Durable)
.build();
let a = MemoryId::content(&owner, MemoryKind::Fact, b"the budget is 5000");
let b = MemoryId::content(&owner, MemoryKind::Fact, b"the budget is 5000");
assert_eq!(a, b, "same owner, kind, and body must dedup to one id");
let other = MemoryId::content(&owner, MemoryKind::Fact, b"the budget is 6000");
assert_ne!(a, other, "a different body must produce a different id");
}
#[test]
fn given_a_known_content_when_addressed_then_should_match_the_cross_sdk_id() {
let owner = MemoryScope::builder()
.agent("agent".parse().expect("valid agent"))
.build();
let id = MemoryId::content(&owner, MemoryKind::Fact, b"x");
assert_eq!(id.to_string(), "1A9GVS6SJ6SNS4KY0H19130WCW");
}
#[tokio::test]
async fn given_positive_feedback_when_recalling_then_should_rank_the_target_first() {
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::builder()
.conversation(ConversationId::new())
.build();
let first = Memory::remember(&memory, &scope, b"the cat sat".to_vec())
.await
.expect("remember first");
let second = Memory::remember(&memory, &scope, b"the dog ran".to_vec())
.await
.expect("remember second");
let before = Memory::recall(&memory, &scope, &MemoryQuery::builder().build())
.await
.expect("recall before feedback");
assert_eq!(before[0].id, first);
Memory::improve(&memory, &scope, Feedback::new(second, 5.0))
.await
.expect("improve");
let after = Memory::recall(&memory, &scope, &MemoryQuery::builder().build())
.await
.expect("recall after feedback");
assert_eq!(after[0].id, second, "the promoted item ranks first");
assert_eq!(after[0].score, Some(5.0));
}
#[test]
fn given_a_named_record_when_absorbed_then_should_fold_into_point_state() {
let mut projection = Projection::default();
let item = laser_wire::memory::MemoryRecord::Item {
id: "notes/current-plan".to_owned(),
kind: "fact".to_owned(),
body: b"ship on friday".to_vec(),
};
projection.absorb(
&laser_wire::framing::encode_named(&item).expect("encodes"),
Provenance::builder()
.conversation_id(ConversationId::new())
.build(),
);
assert_eq!(
projection
.named
.get("notes/current-plan")
.map(Vec::as_slice),
Some(b"ship on friday".as_slice()),
);
assert!(projection.items.is_empty());
let forget = laser_wire::memory::MemoryRecord::Forget {
target: "notes/current-plan".to_owned(),
};
projection.absorb(
&laser_wire::framing::encode_named(&forget).expect("encodes"),
Provenance::builder()
.conversation_id(ConversationId::new())
.build(),
);
assert!(projection.named.is_empty());
}
#[test]
fn given_a_merge_patch_when_applied_then_should_overwrite_add_and_remove_fields() {
let current = br#"{"tier":"pro","seats":3,"beta":true}"#;
let patch = br#"{"tier":"enterprise","region":"eu","beta":null}"#;
let merged = merge_named(Some(current), patch).expect("merges");
let value: serde_json::Value = serde_json::from_slice(&merged).expect("json");
assert_eq!(value["tier"], "enterprise"); assert_eq!(value["seats"], 3); assert_eq!(value["region"], "eu"); assert!(value.get("beta").is_none()); }
#[test]
fn given_no_current_value_when_merge_patched_then_should_start_from_the_patch() {
let merged = merge_named(None, br#"{"plan":"enterprise"}"#).expect("merges");
let value: serde_json::Value = serde_json::from_slice(&merged).expect("json");
assert_eq!(value["plan"], "enterprise");
}
#[tokio::test]
async fn given_a_forgotten_id_when_recalling_then_should_be_excluded() {
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::builder()
.conversation(ConversationId::new())
.build();
let id = Memory::remember(&memory, &scope, b"cat".to_vec())
.await
.expect("remember should succeed");
Memory::forget(&memory, &scope, id)
.await
.expect("forget should succeed");
let all = Memory::recall(&memory, &scope, &MemoryQuery::builder().build())
.await
.expect("recall should succeed");
assert!(all.is_empty());
}
}