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 user = owner.user.as_deref().unwrap_or("").as_bytes();
let application = owner.app.as_deref().unwrap_or("").as_bytes();
let kind = [kind.code()];
let segments: [&[u8]; 10] = [
stream,
&[0],
agent.as_bytes(),
&[0],
user,
&[0],
application,
&[0],
&kind,
body,
];
Self::from_u128(laser_wire::hashing::content_id(&segments))
}
}
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,
pub origin: Option<SourceRef>,
pub producer: Option<laser_wire::graph::ProducerInfo>,
}
#[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>,
pub origin: Option<SourceRef>,
pub producer: Option<laser_wire::graph::ProducerInfo>,
}
impl MemoryItem {
pub fn text(&self) -> String {
String::from_utf8_lossy(&self.payload).into_owned()
}
pub fn json<T: serde::de::DeserializeOwned>(&self) -> Result<T, LaserError> {
serde_json::from_slice(&self.payload).map_err(|error| LaserError::Codec(error.to_string()))
}
}
#[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,
origin: None,
producer: 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 {
usize::try_from(laser_wire::agent::estimate_tokens(text.len())).unwrap_or(usize::MAX)
}
#[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>;
fn append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> impl std::future::Future<Output = Result<MemoryId, LaserError>>
where
Self: Sync,
{
async move {
let _ = (id, kind);
self.remember(scope, payload).await
}
}
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 append<'a>(
&'a self,
scope: &'a MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
> {
let _ = (id, kind);
self.remember(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>,
>;
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 append<'a>(
&'a self,
scope: &'a MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<MemoryId, LaserError>> + Send + 'a>,
> {
Box::pin(Memory::append(self, scope, id, kind, 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>,
#[serde(default, skip_serializing_if = "Option::is_none")]
origin: Option<SourceRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
producer: Option<laser_wire::graph::ProducerInfo>,
},
Forget {
target: MemoryId,
#[serde(default, skip_serializing_if = "Option::is_none")]
conversation: Option<ConversationId>,
},
Feedback {
target: MemoryId,
weight: f32,
#[serde(default, skip_serializing_if = "Option::is_none")]
conversation: Option<ConversationId>,
},
}
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>,
scopes: std::collections::HashMap<MemoryId, StoredMemoryScope>,
forgotten: HashSet<(MemoryId, Option<ConversationId>)>,
feedback: std::collections::HashMap<(MemoryId, Option<ConversationId>), f32>,
named: std::collections::HashMap<String, Vec<u8>>,
offsets: Vec<u64>,
}
#[derive(Default)]
struct StoredMemoryScope {
user: Option<String>,
app: Option<String>,
}
impl Projection {
#[cfg(test)]
fn absorb(&mut self, payload: &[u8], provenance: Provenance) {
self.absorb_scoped(payload, provenance, StoredMemoryScope::default(), None);
}
fn absorb_scoped(
&mut self,
payload: &[u8],
provenance: Provenance,
scope: StoredMemoryScope,
source: Option<SourceRef>,
) {
let Ok(entry) = laser_wire::framing::decode_named::<MemoryLogEntry>(payload) else {
self.absorb_named(payload);
return;
};
match entry {
MemoryLogEntry::Forget {
target,
conversation,
} => {
let before = self.items.len();
self.items.retain(|item| {
item.id != target
|| conversation
.is_some_and(|scope| scope != item.provenance.conversation_id)
});
if self.items.len() != before {
self.scopes.remove(&target);
}
self.forgotten.insert((target, conversation));
}
MemoryLogEntry::Feedback {
target,
weight,
conversation,
} => {
*self.feedback.entry((target, conversation)).or_insert(0.0) += weight;
}
MemoryLogEntry::Item {
id,
kind,
body,
origin,
producer,
} => {
if self.is_forgotten(id, provenance.conversation_id) {
return;
}
if self.items.iter().any(|item| item.id == id) {
return;
}
self.scopes.insert(id, scope);
self.items.push(MemoryItem {
id,
payload: body,
provenance,
source,
kind,
score: None,
signals: Vec::new(),
origin,
producer,
});
}
}
}
fn is_forgotten(&self, id: MemoryId, conversation: ConversationId) -> bool {
self.forgotten.contains(&(id, None)) || self.forgotten.contains(&(id, Some(conversation)))
}
fn feedback_for(&self, item: &MemoryItem) -> Option<f32> {
let any = self.feedback.get(&(item.id, None));
let scoped = self
.feedback
.get(&(item.id, Some(item.provenance.conversation_id)));
match (any, scoped) {
(None, None) => None,
(any, scoped) => Some(any.copied().unwrap_or(0.0) + scoped.copied().unwrap_or(0.0)),
}
}
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 { .. } => {}
}
}
}
fn folded_memory_items(
projection: &Projection,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Vec<MemoryItem> {
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.is_forgotten(item.id, item.provenance.conversation_id))
.filter(|item| {
let stored = projection.scopes.get(&item.id);
scope.user.as_ref().is_none_or(|expected| {
stored.and_then(|value| value.user.as_ref()) == Some(expected)
}) && scope.app.as_ref().is_none_or(|expected| {
stored.and_then(|value| value.app.as_ref()) == Some(expected)
})
})
.filter(|item| {
agent_filter.is_none_or(|agent| item.provenance.agent.as_ref() == Some(agent))
})
.filter(|item| memory_query_matches(query, &item.payload))
.cloned()
.collect();
items.reverse();
if query.strategy != RecallStrategy::Recent && !projection.feedback.is_empty() {
for item in &mut items {
if let Some(weight) = projection.feedback_for(item) {
item.score = Some(weight);
}
}
items.sort_by(|left, right| {
right
.score
.unwrap_or(0.0)
.total_cmp(&left.score.unwrap_or(0.0))
});
for (rank, item) in items.iter_mut().enumerate() {
if let Some(score) = item.score {
item.signals.push(RecallSignal {
strategy: RecallStrategy::Auto,
rank,
score: Some(score),
});
}
}
}
items.truncate(query.limit);
items
}
impl LogMemory {
fn stream_name(&self) -> Result<&str, LaserError> {
match self.stream.as_deref() {
Some(stream) => Ok(stream),
None => self.laser.stream_required(),
}
}
fn resource_namespace(&self) -> String {
self.laser
.resource_name_in(self.resource_stream(), &self.namespace)
}
fn resource_stream(&self) -> Option<&str> {
self.laser
.resource_scope(self.stream.as_deref().or(self.laser.default_stream()))
}
pub fn new(laser: Laser) -> Self {
Self {
laser,
stream: None,
topic: AgentTopic::Memory.as_identifier(),
namespace: AgentTopic::Memory.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::Memory.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 action = crate::govern::GovernedAction {
kind: crate::govern::ActionKind::MemoryWrite,
stream: self.stream_name()?,
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.resource_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,
origin: scope.origin.clone(),
producer: scope.producer.clone(),
};
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 {
origin: None,
producer: None,
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(),
conversation: None,
};
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.resource_namespace())
.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 = Identifier::named(self.stream_name()?)?;
let topic = self.topic.clone();
let Some(details) = self.laser.client().get_topic(&stream, &topic).await? else {
return Ok(());
};
let stream_id = self
.laser
.client()
.get_stream(&stream)
.await?
.ok_or_else(|| LaserError::Invalid("memory stream is missing".to_owned()))?
.id;
let partitions = crate::poll::bounded_partitions(details.partitions_count);
let consumer = Consumer::new(Identifier::named("laser-log-memory")?);
let namespace = self.resource_namespace();
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,
crate::poll::DrainRange::open(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(Some(headers)) = message.user_headers_map() else {
continue;
};
let Ok(provenance) = crate::provenance::provenance_from_headers(&headers) else {
continue;
};
let text = |name: &str| -> Option<String> {
let key = HeaderKey::from_str(name).ok()?;
headers.get(&key)?.as_str().ok().map(str::to_owned)
};
if text(laser_wire::headers::MEMORY_NAMESPACE).as_deref()
!= Some(namespace.as_str())
{
continue;
}
let source = SourceRef::Message {
generation: Some(details.created_at.as_micros()),
stream: stream_id,
topic: details.id,
partition,
offset: message.header.offset,
conversation: Some(provenance.conversation_id.to_string()),
};
projection.absorb_scoped(
&message.payload,
provenance,
StoredMemoryScope {
user: text(laser_wire::headers::MEMORY_USER),
app: text(laser_wire::headers::MEMORY_APP),
},
Some(source),
);
}
projection.offsets[partition as usize] = next_offset;
}
Ok(())
}
pub async fn recall_folded(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
if scope
.stream
.as_deref()
.is_some_and(|stream| self.stream_name().ok() != Some(stream))
{
return Ok(Vec::new());
}
self.catch_up().await?;
let projection = self.projection.lock().await;
Ok(folded_memory_items(&projection, scope, query))
}
#[cfg(feature = "kv")]
async fn recall_kv(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
if scope
.stream
.as_deref()
.is_some_and(|stream| self.stream_name().ok() != Some(stream))
{
return Ok(Vec::new());
}
let kv = self.laser.kv(self.resource_namespace());
let lens_stream = self.resource_stream().map(str::to_owned);
recall_kv_pages(scope, query, |cursor| {
let mut scan = kv.scan();
if let Some(conversation) = scope.conversation {
scan = scan
.conversation(conversation)
.lens_stream(lens_stream.clone());
}
if let Some(cursor) = cursor {
scan = scan.cursor(cursor);
}
async move { scan.fetch().await }
})
.await
}
}
#[cfg(feature = "kv")]
async fn recall_kv_pages<F, Fut>(
scope: &MemoryScope,
query: &MemoryQuery,
mut fetch: F,
) -> Result<Vec<MemoryItem>, LaserError>
where
F: FnMut(Option<Vec<u8>>) -> Fut,
Fut: std::future::Future<Output = Result<laser_wire::kv::KvPage, LaserError>>,
{
if query.limit == 0 {
return Ok(Vec::new());
}
let mut selected = std::collections::BTreeMap::new();
let mut cursor = None;
loop {
let page = fetch(cursor.clone()).await?;
for entry in page.entries {
let arrival = arrival_key(&entry);
if let Some(item) = memory_item_from_entry(scope, query.agent.as_ref(), entry)
.filter(|item| memory_query_matches(query, &item.payload))
{
selected.insert((arrival, item.id), item);
if selected.len() > query.limit {
selected.pop_first();
}
}
}
let Some(next) = page.cursor else { break };
if cursor.as_ref() == Some(&next) {
return Err(LaserError::Protocol(
"memory scan cursor did not advance".to_owned(),
));
}
cursor = Some(next);
}
Ok(selected.into_values().rev().collect())
}
#[cfg(feature = "kv")]
fn arrival_key(entry: &laser_wire::kv::KvEntry) -> (u64, u32, u32, u64) {
let scope = entry.scope.as_deref();
let timestamp = scope.and_then(|scope| scope.timestamp_micros).unwrap_or(0);
match scope.and_then(|scope| scope.source.as_ref()) {
Some(SourceRef::Message {
topic,
partition,
offset,
..
}) => (timestamp, *topic, *partition, *offset),
_ => (timestamp, 0, 0, 0),
}
}
#[cfg(feature = "kv")]
fn memory_item_from_entry(
scope: &MemoryScope,
agent: Option<&AgentId>,
entry: laser_wire::kv::KvEntry,
) -> Option<MemoryItem> {
let id: MemoryId = entry.key_str()?.parse().ok()?;
let stored = *entry.scope?;
let conversation = match stored.conversation.as_deref() {
Some(value) => value.parse().ok()?,
None => scope.conversation?,
};
let agent = agent.or(scope.agent.as_ref());
if scope
.conversation
.is_some_and(|expected| expected != conversation)
|| scope
.user
.as_ref()
.is_some_and(|expected| stored.user.as_ref() != Some(expected))
|| scope
.app
.as_ref()
.is_some_and(|expected| stored.app.as_ref() != Some(expected))
|| agent.is_some_and(|expected| stored.agent.as_deref() != Some(expected.as_str()))
{
return None;
}
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(|value| value.parse().ok());
Some(MemoryItem {
id,
payload: entry.value,
provenance,
kind,
source: stored.source,
score: None,
signals: Vec::new(),
origin: None,
producer: None,
})
}
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 append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
LogMemory::append(self, scope, id, kind, 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,
conversation: scope.conversation,
};
let provenance = Self::provenance(scope, id.to_string());
self.publish(&provenance, scope, entry.encode()?).await?;
Ok(id)
}
async fn forget(&self, scope: &MemoryScope, id: MemoryId) -> Result<(), LaserError> {
let entry = MemoryLogEntry::Forget {
target: id,
conversation: scope.conversation,
};
let provenance = Self::provenance(scope, MemoryId::new().to_string());
self.publish(&provenance, scope, entry.encode()?).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)
.strategy(RecallStrategy::Recent)
.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() {
for (conversation, turns) in by_conversation(&sessions) {
let bodies = turns.iter().map(|item| item.payload.clone()).collect();
let summary = summarizer.summarize(bodies).await?;
let mut summary_scope = scope.clone();
summary_scope.conversation = Some(conversation);
summary_scope.lifetime = Lifetime::Durable;
self.memory
.append(
&summary_scope,
MemoryId::new(),
MemoryKind::Summary,
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);
}
let prune_count = items.len() - self.max_items;
for item in items.into_iter().rev().take(prune_count) {
if self.memory.forget(scope, item.id).await.is_ok() {
report.pruned += 1;
}
}
Ok(report)
}
}
fn by_conversation<'a>(items: &[&'a MemoryItem]) -> Vec<(ConversationId, Vec<&'a MemoryItem>)> {
let mut groups: Vec<(ConversationId, Vec<&'a MemoryItem>)> = Vec::new();
for &item in items {
let conversation = item.provenance.conversation_id;
match groups.iter_mut().find(|(known, _)| *known == conversation) {
Some((_, turns)) => turns.push(item),
None => groups.push((conversation, vec![item])),
}
}
groups
}
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 append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.inner.append(scope, id, kind, 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::Memory.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,
origin,
producer,
..
} => Ok(MemoryLogEntry::Item {
id,
kind,
body: modified,
origin,
producer,
}),
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,
origin: scope.origin.clone(),
producer: scope.producer.clone(),
},
)
.await?;
let MemoryLogEntry::Item {
id,
kind,
body,
origin,
producer,
} = 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,
origin,
producer,
};
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 append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
VectorMemory::append(self, scope, id, kind, payload).await
}
async fn recall(
&self,
scope: &MemoryScope,
query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
let wants_semantic = !matches!(
query.strategy,
RecallStrategy::Keyword | RecallStrategy::Recent
);
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))
&& scope
.stream
.as_ref()
.is_none_or(|stream| entry.scope.stream.as_ref() == Some(stream))
})
.collect();
let feedback_ranks = query.strategy != RecallStrategy::Recent;
let any_feedback = feedback_ranks && matched.iter().any(|entry| entry.feedback != 0.0);
let mut scored: Vec<(f32, &VectorEntry, Vec<RecallSignal>)> = matched
.iter()
.map(|&entry| {
let mut signals = Vec::new();
if let Some(embedding) = &query_embedding {
signals.push(RecallSignal {
strategy: RecallStrategy::Semantic,
rank: 0,
score: Some(cosine(embedding, &entry.embedding)),
});
}
if let Some(tokens) = &query_tokens {
signals.push(RecallSignal {
strategy: RecallStrategy::Keyword,
rank: 0,
score: Some(keyword_score(tokens, &entry.item.payload)),
});
}
if feedback_ranks && entry.feedback != 0.0 {
signals.push(RecallSignal {
strategy: RecallStrategy::Auto,
rank: 0,
score: Some(entry.feedback),
});
}
let score = signals
.iter()
.map(|signal| signal.score.unwrap_or(0.0))
.sum();
(score, entry, signals)
})
.collect();
let any_query_signal = query_embedding.is_some() || query_tokens.is_some();
if any_query_signal || any_feedback {
for strategy in [
RecallStrategy::Semantic,
RecallStrategy::Keyword,
RecallStrategy::Auto,
] {
let mut ranks: Vec<(usize, f32)> = scored
.iter()
.enumerate()
.filter_map(|(index, (_, _, signals))| {
signals
.iter()
.find(|signal| signal.strategy == strategy)
.map(|signal| (index, signal.score.unwrap_or(0.0)))
})
.collect();
ranks.sort_by(|left, right| right.1.total_cmp(&left.1));
for (rank, (index, _)) in ranks.into_iter().enumerate() {
if let Some(signal) = scored[index]
.2
.iter_mut()
.find(|signal| signal.strategy == strategy)
{
signal.rank = rank;
}
}
}
scored.sort_by(|left, right| right.0.total_cmp(&left.0));
Ok(scored
.into_iter()
.take(query.limit)
.map(|(score, entry, signals)| {
let mut item = entry.item.clone();
item.score = Some(score);
item.signals = signals;
item
})
.collect())
} else {
Ok(matched
.into_iter()
.rev()
.take(query.limit)
.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,
conversation: scope.conversation,
},
)
.await?;
let MemoryLogEntry::Feedback {
target,
weight,
conversation,
} = 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 && in_conversation(entry, conversation))
{
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,
conversation: scope.conversation,
},
)
.await?;
let MemoryLogEntry::Forget {
target,
conversation,
} = 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 || !in_conversation(entry, conversation));
Ok(())
}
}
fn in_conversation(entry: &VectorEntry, conversation: Option<ConversationId>) -> bool {
conversation.is_none_or(|scope| entry.item.provenance.conversation_id == scope)
}
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_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
}
fn memory_query_matches(query: &MemoryQuery, body: &[u8]) -> bool {
if query.strategy == RecallStrategy::Recent {
return true;
}
let Some(text) = query.semantic.as_deref() else {
return true;
};
let tokens = tokenize(text);
tokens.is_empty() || keyword_score(&tokens, body) > 0.0
}
#[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",
}
}
pub fn vector(embedder: impl Embedder + Sync + 'static) -> Self {
Self::Vector(std::sync::Arc::new(VectorMemory::new(SharedEmbedder(
std::sync::Arc::new(embedder),
))))
}
#[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>,
stream: Option<&str>,
) -> Result<MemoryHandle, LaserError> {
Ok(MemoryHandle::Log(LogMemory::on_stream_topic_named(
self.clone(),
stream.map(str::to_owned),
topic.as_ref(),
)?))
}
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.append(scope, id, kind, 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, reranker } => {
let items = recall_folded_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,
}
}
pub fn recall(&self, conversation: impl Into<Option<ConversationId>>) -> RecallBuilder<'_> {
RecallBuilder {
handle: self,
scope: MemoryScope::builder()
.maybe_conversation(conversation.into())
.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 consolidate_with(
&self,
scope: &MemoryScope,
max_items: usize,
summarizer: impl Summarizer + Sync,
prune_summarized: bool,
) -> Result<ConsolidationReport, LaserError> {
let consolidator = DefaultConsolidator::new(self, max_items).with_summarizer(summarizer);
let consolidator = if prune_summarized {
consolidator.prune_summarized()
} else {
consolidator
};
Consolidator::consolidate(&consolidator, 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 append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.append_id(scope, id, kind, 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
}
#[allow(clippy::manual_async_fn)]
fn append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> impl std::future::Future<Output = Result<MemoryId, LaserError>> + Send
where
Self: Sync,
{
async move { Memory::append(*self, scope, id, kind, 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 origin(mut self, origin: SourceRef) -> Self {
self.scope.origin = Some(origin);
self
}
#[must_use]
pub fn producer(mut self, producer: laser_wire::graph::ProducerInfo) -> Self {
self.scope.producer = Some(producer);
self
}
#[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 user(mut self, user: impl Into<String>) -> Self {
self.scope.user = Some(user.into());
self
}
#[must_use]
pub fn application(mut self, application: impl Into<String>) -> Self {
self.scope.app = Some(application.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 user(mut self, user: impl Into<String>) -> Self {
self.scope.user = Some(user.into());
self
}
#[must_use]
pub fn application(mut self, application: impl Into<String>) -> Self {
self.scope.app = Some(application.into());
self
}
#[must_use]
pub fn stream(mut self, stream: impl Into<String>) -> Self {
self.scope.stream = Some(stream.into());
self
}
#[must_use]
pub fn limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
pub async fn block(self, token_budget: Option<usize>) -> Result<String, LaserError> {
Ok(to_context_block(&self.fetch().await?, token_budget))
}
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
}
}
}
impl<'a> std::future::IntoFuture for RecallBuilder<'a> {
type Output = Result<Vec<MemoryItem>, LaserError>;
type IntoFuture =
std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send + 'a>>;
fn into_future(self) -> Self::IntoFuture {
Box::pin(self.fetch())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn given_a_default_stream_when_naming_the_memory_namespace_then_should_scope_it_to_the_record_stream()
{
let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
.with_default_stream("acme");
let memory = LogMemory::in_namespace(laser.clone(), "notes");
assert_eq!(memory.resource_namespace(), "stream:acme/notes");
assert_eq!(memory.named_key("k"), "notes/k");
let elsewhere =
LogMemory::on_stream_topic_named(laser.clone(), Some("fleet".to_owned()), "memory")
.expect("a valid topic name");
assert_eq!(elsewhere.resource_namespace(), "stream:fleet/memory");
assert_eq!(elsewhere.resource_stream(), Some("fleet"));
let bare = LogMemory::in_namespace(
laser.with_resource_naming(crate::laser::ResourceNaming::Bare),
"notes",
);
assert_eq!(bare.resource_namespace(), "notes");
assert_eq!(bare.resource_stream(), None);
}
#[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");
}
#[tokio::test]
async fn given_user_and_application_when_building_then_should_reach_the_backend_scope() {
type Seen = std::sync::Arc<std::sync::Mutex<Vec<MemoryScope>>>;
struct ScopeProbe(Seen);
impl Memory for ScopeProbe {
async fn remember(
&self,
scope: &MemoryScope,
_payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.0.lock().expect("lock").push(scope.clone());
Ok(MemoryId::new())
}
async fn recall(
&self,
scope: &MemoryScope,
_query: &MemoryQuery,
) -> Result<Vec<MemoryItem>, LaserError> {
self.0.lock().expect("lock").push(scope.clone());
let provenance = Provenance::builder()
.conversation_id(ConversationId::new())
.build();
Ok(vec![MemoryItem::plain(
MemoryId::new(),
b"node-7 rotates keys monthly".to_vec(),
provenance,
)])
}
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 seen = Seen::default();
let handle = MemoryHandle::Custom(std::sync::Arc::new(ScopeProbe(seen.clone())));
handle
.remember(b"fact".to_vec())
.user("operator-1")
.application("fleet")
.stream("fleet-memory")
.send()
.await
.expect("the scoped remember should succeed");
let block = handle
.recall(ConversationId::new())
.user("operator-1")
.application("fleet")
.stream("fleet-memory")
.block(Some(100))
.await
.expect("the block should render");
assert_eq!(block, "node-7 rotates keys monthly");
let seen = seen.lock().expect("lock");
assert_eq!(seen.len(), 2);
for scope in seen.iter() {
assert_eq!(scope.user.as_deref(), Some("operator-1"));
assert_eq!(scope.app.as_deref(), Some("fleet"));
assert_eq!(scope.stream.as_deref(), Some("fleet-memory"));
}
}
#[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"auth is slow".to_vec(),
origin: None,
producer: None,
};
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 {
origin: None,
producer: None,
id: rid,
kind,
body,
} => {
assert_eq!(rid, id.to_string());
assert_eq!(kind, "message");
assert_eq!(body, b"auth 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()
.rev()
.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);
});
}
#[test]
fn given_content_ids_when_consolidated_then_should_prune_in_arrival_order() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.expect("runtime builds");
runtime.block_on(async {
let scope = MemoryScope::builder().build();
let candidates: Vec<Vec<u8>> = (0..100)
.map(|number| format!("item-{number}").into_bytes())
.collect();
let first = candidates
.iter()
.max_by_key(|body| MemoryId::content(&scope, MemoryKind::Fact, body))
.expect("candidates")
.clone();
let first_id = MemoryId::content(&scope, MemoryKind::Fact, &first);
let newest = candidates
.iter()
.min_by_key(|body| MemoryId::content(&scope, MemoryKind::Fact, body))
.expect("candidates")
.clone();
let newest_id = MemoryId::content(&scope, MemoryKind::Fact, &newest);
assert!(newest_id < first_id);
let provenance = Provenance::builder()
.conversation_id(crate::types::ConversationId::new())
.build();
let mock = MockMemory {
items: std::sync::Mutex::new(vec![
MemoryItem::plain(first_id, first, provenance.clone()),
MemoryItem::plain(newest_id, newest.clone(), provenance),
]),
};
let consolidator = DefaultConsolidator::new(mock, 1);
let report = Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidate");
assert_eq!(report.pruned, 1);
let remaining = Memory::recall(
&consolidator.memory,
&scope,
&MemoryQuery::builder().limit(10).build(),
)
.await
.expect("recall");
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].payload, newest);
});
}
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();
let conversation = ConversationId::new();
for body in [b"turn one".to_vec(), b"turn two".to_vec()] {
memory.push_in(MemoryKind::Message, body, conversation);
}
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_turns_in_two_conversations_when_consolidated_then_should_store_one_summary_each()
{
let memory = LogMemoryStub::default();
let first = ConversationId::new();
let second = ConversationId::new();
memory.push_in(MemoryKind::Message, b"first a".to_vec(), first);
memory.push_in(MemoryKind::Message, b"second a".to_vec(), second);
memory.push_in(MemoryKind::Message, b"first b".to_vec(), first);
let scope = MemoryScope::builder()
.agent("planner".parse().expect("valid agent"))
.build();
let consolidator =
DefaultConsolidator::new(memory.clone(), 100).with_summarizer(StubSummarizer);
let report = Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidates");
assert_eq!(report.summarized, 3);
let summaries: Vec<_> = memory
.appended()
.into_iter()
.map(|(scope, kind, body)| {
(scope.conversation, scope.agent, scope.lifetime, kind, body)
})
.collect();
let planner = Some("planner".parse().expect("valid agent"));
assert_eq!(
summaries,
[
(
Some(first),
planner.clone(),
Lifetime::Durable,
MemoryKind::Summary,
b"summary of 2".to_vec()
),
(
Some(second),
planner,
Lifetime::Durable,
MemoryKind::Summary,
b"summary of 1".to_vec()
),
]
);
let stored = memory.items.lock().expect("stub lock").clone();
let summary_conversations: Vec<_> = stored
.iter()
.filter(|item| item.kind == MemoryKind::Summary)
.map(|item| item.provenance.conversation_id)
.collect();
assert_eq!(summary_conversations, [first, second]);
}
#[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");
}
type Appended = (MemoryScope, MemoryKind, Vec<u8>);
#[derive(Clone, Default)]
struct LogMemoryStub {
items: std::sync::Arc<std::sync::Mutex<Vec<MemoryItem>>>,
appended: std::sync::Arc<std::sync::Mutex<Vec<Appended>>>,
}
impl LogMemoryStub {
fn push(&self, kind: MemoryKind, payload: Vec<u8>) {
self.push_in(kind, payload, crate::types::ConversationId::new());
}
fn push_in(&self, kind: MemoryKind, payload: Vec<u8>, conversation: ConversationId) {
let provenance = Provenance::builder().conversation_id(conversation).build();
let mut item = MemoryItem::plain(MemoryId::new(), payload, provenance);
item.kind = kind;
self.items.lock().expect("stub lock").push(item);
}
fn appended(&self) -> Vec<Appended> {
self.appended.lock().expect("stub lock").clone()
}
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 append(
&self,
scope: &MemoryScope,
id: MemoryId,
kind: MemoryKind,
payload: Vec<u8>,
) -> Result<MemoryId, LaserError> {
self.appended
.lock()
.expect("stub lock")
.push((scope.clone(), kind, payload.clone()));
self.push_in(kind, payload, scope.conversation.unwrap_or_default());
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_a_text_payload_when_read_then_should_return_the_text_and_json() {
let item = context_item(r#"{"city":"Oslo"}"#);
assert_eq!(item.text(), r#"{"city":"Oslo"}"#);
let value: serde_json::Value = item.json().expect("json");
assert_eq!(value["city"], "Oslo");
}
#[test]
fn given_a_non_json_payload_when_decoded_then_should_fail_with_a_codec_error() {
let item = context_item("plain words");
assert!(matches!(
item.json::<serde_json::Value>(),
Err(LaserError::Codec(_))
));
}
#[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(),
origin: None,
producer: None,
}
.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,
conversation: None,
}
.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.is_forgotten(id_a, conversation),
"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,
conversation: None,
}
.encode()
.expect("a feedback entry encodes");
projection.absorb(&feedback, prov());
projection.absorb(&feedback, prov());
assert_eq!(projection.feedback.get(&(target, None)), Some(&5.0));
assert!(projection.items.is_empty(), "feedback is not an item");
}
#[test]
fn given_scoped_tombstone_and_feedback_when_absorbed_then_should_touch_only_their_conversation()
{
let owner = ConversationId::new();
let other = ConversationId::new();
let in_owner = || Provenance::builder().conversation_id(owner).build();
let in_other = || Provenance::builder().conversation_id(other).build();
let mut projection = Projection::default();
let id = MemoryId::new();
projection.absorb(&item_entry(id, b"kept"), in_owner());
let scoped = |entry: MemoryLogEntry| entry.encode().expect("entry encodes");
projection.absorb(
&scoped(MemoryLogEntry::Feedback {
target: id,
weight: 4.0,
conversation: Some(other),
}),
in_other(),
);
projection.absorb(
&scoped(MemoryLogEntry::Forget {
target: id,
conversation: Some(other),
}),
in_other(),
);
assert_eq!(
projection.items.len(),
1,
"another conversation cannot forget it"
);
assert_eq!(projection.feedback_for(&projection.items[0]), None);
projection.absorb(
&scoped(MemoryLogEntry::Feedback {
target: id,
weight: 2.0,
conversation: Some(owner),
}),
in_owner(),
);
assert_eq!(projection.feedback_for(&projection.items[0]), Some(2.0));
projection.absorb(
&scoped(MemoryLogEntry::Forget {
target: id,
conversation: Some(owner),
}),
in_owner(),
);
assert!(
projection.items.is_empty(),
"the owning conversation forgets it"
);
}
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"
);
}
#[tokio::test]
async fn given_a_reranked_handle_when_recalling_folded_then_should_apply_the_reranker() {
let handle = MemoryHandle::vector(WordEmbedder).reranker(ReverseReranker);
let conversation = ConversationId::new();
for body in ["the cat sat", "the dog ran"] {
handle
.remember(body)
.scope(conversation)
.send()
.await
.expect("remember");
}
let scope = MemoryScope::builder().conversation(conversation).build();
let query = MemoryQuery::builder().semantic("cat".to_owned()).build();
let hits = handle
.recall_folded(&scope, &query)
.await
.expect("folded recall");
assert_eq!(
hits.iter()
.map(|item| item.payload.as_slice())
.collect::<Vec<_>>(),
[b"the dog ran".as_slice(), b"the cat sat"],
"the reranker reversed the similarity order"
);
}
#[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 agent_only = MemoryScope::builder()
.agent("agent".parse().expect("valid agent"))
.build();
let full_owner = MemoryScope::builder()
.stream("laser".to_owned())
.agent("agent".parse().expect("valid agent"))
.user("alice".to_owned())
.app("console".to_owned())
.build();
let user_only = MemoryScope::builder().user("alice".to_owned()).build();
let cases = [
(
&agent_only,
MemoryKind::Fact,
"63FZWE3WTSCVXAY845QK4TMVWM",
0xc37f_f8e1_f359_66fa_af20_85bc_c9aa_6f94_u128,
),
(
&full_owner,
MemoryKind::Fact,
"28429K3B8MPSM8WHDQR474827P",
0x4820_9331_ad14_b668_8e45_b7c1_0e44_08f6,
),
(
&user_only,
MemoryKind::Procedure,
"3E6HJGHKFDBMQ91TCJEHV28X40",
0x6e34_6508_cded_5d2e_90e9_9274_7624_7480,
),
];
for (owner, kind, rendered, value) in cases {
let id = MemoryId::content(owner, kind, b"x");
assert_eq!(id.as_u128(), value);
assert_eq!(id.to_string(), rendered);
}
}
#[test]
fn given_two_users_when_the_same_body_is_addressed_then_should_produce_distinct_ids() {
let owner = |user: &str, application: &str| {
MemoryScope::builder()
.agent("agent".parse().expect("valid agent"))
.user(user.to_owned())
.app(application.to_owned())
.build()
};
let alice = MemoryId::content(&owner("alice", "console"), MemoryKind::Fact, b"x");
let bob = MemoryId::content(&owner("bob", "console"), MemoryKind::Fact, b"x");
let support = MemoryId::content(&owner("alice", "support"), MemoryKind::Fact, b"x");
assert_ne!(alice, bob, "the user is part of the owner");
assert_ne!(alice, support, "the application is part of the owner");
let shifted = MemoryId::content(&owner("alicecon", "sole"), MemoryKind::Fact, b"x");
assert_ne!(
owner("alicecon", "sole").user,
owner("alice", "console").user
);
assert_ne!(alice, shifted);
}
#[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, second);
Memory::improve(&memory, &scope, Feedback::new(first, 5.0))
.await
.expect("improve");
let after = Memory::recall(&memory, &scope, &MemoryQuery::builder().build())
.await
.expect("recall after feedback");
assert_eq!(after[0].id, first, "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 {
origin: None,
producer: None,
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(),
conversation: None,
};
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#"{"log_level":"info","replicas":3,"beta":true}"#;
let patch = br#"{"log_level":"debug","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["log_level"], "debug"); assert_eq!(value["replicas"], 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#"{"log_level":"debug"}"#).expect("merges");
let value: serde_json::Value = serde_json::from_slice(&merged).expect("json");
assert_eq!(value["log_level"], "debug");
}
#[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());
}
#[cfg(feature = "kv")]
fn scoped_kv_entry(
id: u128,
conversation: ConversationId,
user: &str,
app: &str,
) -> laser_wire::kv::KvEntry {
laser_wire::kv::KvEntry {
key: Ulid::from(id).to_string().into_bytes(),
value: format!("item-{id}").into_bytes(),
expires_at_micros: None,
version: 1,
scope: Some(Box::new(laser_wire::kv::MemoryRowScope {
kind: Some("fact".to_owned()),
user: Some(user.to_owned()),
app: Some(app.to_owned()),
conversation: Some(conversation.to_string()),
..Default::default()
})),
source: None,
}
}
#[cfg(feature = "kv")]
#[tokio::test]
async fn given_paged_memory_scopes_when_recalled_without_a_conversation_then_should_keep_only_newest_matching_items()
{
let first = ConversationId::new();
let second = ConversationId::new();
let scope = MemoryScope::builder()
.user("reader".to_owned())
.app("diagnostics".to_owned())
.build();
let query = MemoryQuery::builder().limit(2).build();
let mut pages = std::collections::VecDeque::from([
laser_wire::kv::KvPage {
entries: vec![
scoped_kv_entry(1, first, "reader", "diagnostics"),
scoped_kv_entry(2, first, "other", "diagnostics"),
],
cursor: Some(vec![2]),
},
laser_wire::kv::KvPage {
entries: vec![
scoped_kv_entry(3, second, "reader", "other"),
scoped_kv_entry(4, second, "reader", "diagnostics"),
scoped_kv_entry(5, first, "reader", "diagnostics"),
],
cursor: None,
},
]);
let mut requested = Vec::new();
let items = recall_kv_pages(&scope, &query, |cursor| {
requested.push(cursor);
std::future::ready(Ok(pages.pop_front().expect("next page")))
})
.await
.expect("recall");
assert_eq!(requested, [None, Some(vec![2])]);
assert_eq!(
items
.iter()
.map(|item| item.id.to_string())
.collect::<Vec<_>>(),
[Ulid::from(5u128).to_string(), Ulid::from(4u128).to_string()]
);
assert_eq!(items[0].provenance.conversation_id, first);
assert_eq!(items[1].provenance.conversation_id, second);
}
#[cfg(feature = "kv")]
#[tokio::test]
async fn given_content_ids_out_of_arrival_order_when_recalled_from_the_view_then_should_keep_the_latest_broker_time()
{
let conversation = ConversationId::new();
let stamped = |id: u128, partition: u32, timestamp: u64| {
let mut entry = scoped_kv_entry(id, conversation, "reader", "diagnostics");
let scope = entry.scope.as_mut().expect("memory scope");
scope.timestamp_micros = Some(timestamp);
scope.source = Some(SourceRef::Message {
stream: 1,
topic: 2,
partition,
offset: 0,
generation: None,
conversation: None,
});
entry
};
let scope = MemoryScope::default();
let query = MemoryQuery::builder().limit(2).build();
let items = recall_kv_pages(&scope, &query, |_| {
std::future::ready(Ok(laser_wire::kv::KvPage {
entries: vec![stamped(9, 0, 10), stamped(1, 1, 30), stamped(5, 2, 20)],
cursor: None,
}))
})
.await
.expect("recall");
assert_eq!(
items
.iter()
.map(|item| item.id.to_string())
.collect::<Vec<_>>(),
[Ulid::from(1u128).to_string(), Ulid::from(5u128).to_string()]
);
}
#[cfg(feature = "kv")]
#[tokio::test]
async fn given_a_repeating_memory_scan_cursor_when_recalled_then_should_refuse_to_loop() {
let scope = MemoryScope::default();
let query = MemoryQuery::builder().build();
let result = recall_kv_pages(&scope, &query, |_| {
std::future::ready(Ok(laser_wire::kv::KvPage {
entries: Vec::new(),
cursor: Some(vec![1]),
}))
})
.await;
assert!(matches!(result, Err(LaserError::Protocol(_))));
}
#[test]
fn given_folded_user_and_application_scopes_when_recalled_then_should_keep_the_newest_matching_records()
{
let conversation = ConversationId::new();
let mut projection = Projection::default();
for (id, user, app) in [
(1u128, "reader", "diagnostics"),
(2, "other", "diagnostics"),
(3, "reader", "other"),
(4, "reader", "diagnostics"),
] {
let item = MemoryLogEntry::Item {
id: MemoryId(Ulid::from(id)),
kind: MemoryKind::Fact,
body: format!("item-{id}").into_bytes(),
origin: None,
producer: None,
};
projection.absorb_scoped(
&item.encode().expect("encode"),
Provenance::builder().conversation_id(conversation).build(),
StoredMemoryScope {
user: Some(user.to_owned()),
app: Some(app.to_owned()),
},
None,
);
}
let scope = MemoryScope::builder()
.conversation(conversation)
.user("reader".to_owned())
.app("diagnostics".to_owned())
.build();
let items = folded_memory_items(
&projection,
&scope,
&MemoryQuery::builder().limit(1).build(),
);
assert_eq!(items.len(), 1);
assert_eq!(items[0].payload, b"item-4");
}
#[test]
fn given_folded_log_memory_when_searched_then_should_exclude_unrelated_items_before_limiting() {
let conversation = ConversationId::new();
let mut projection = Projection::default();
for (id, body) in [(1u128, b"apple".as_slice()), (2, b"orange"), (3, b"pear")] {
projection.absorb(
&item_entry(MemoryId(Ulid::from(id)), body),
Provenance::builder().conversation_id(conversation).build(),
);
}
for strategy in [
RecallStrategy::Keyword,
RecallStrategy::Semantic,
RecallStrategy::Hybrid,
] {
let query = MemoryQuery::builder()
.semantic("apple".to_owned())
.strategy(strategy)
.limit(1)
.build();
let items = folded_memory_items(&projection, &MemoryScope::default(), &query);
assert_eq!(items.len(), 1);
assert_eq!(items[0].payload, b"apple");
}
}
#[cfg(feature = "kv")]
#[tokio::test]
async fn given_managed_log_memory_when_searched_then_should_scan_past_unrelated_newer_items() {
let conversation = ConversationId::new();
let mut matching = scoped_kv_entry(1, conversation, "reader", "diagnostics");
matching.value = b"apple".to_vec();
let mut unrelated = scoped_kv_entry(2, conversation, "reader", "diagnostics");
unrelated.value = b"orange".to_vec();
let scope = MemoryScope::builder().conversation(conversation).build();
let query = MemoryQuery::builder()
.semantic("apple".to_owned())
.strategy(RecallStrategy::Keyword)
.limit(1)
.build();
let mut page = Some(laser_wire::kv::KvPage {
entries: vec![matching, unrelated],
cursor: None,
});
let items = recall_kv_pages(&scope, &query, |_| {
std::future::ready(Ok(page.take().expect("one page")))
})
.await
.expect("recall");
assert_eq!(items.len(), 1);
assert_eq!(items[0].payload, b"apple");
}
#[test]
fn given_a_folded_memory_record_when_recalled_then_should_keep_its_source() {
let conversation = ConversationId::new();
let source = SourceRef::Message {
generation: None,
stream: 1,
topic: 2,
partition: 0,
offset: 9,
conversation: Some(conversation.to_string()),
};
let entry = MemoryLogEntry::Item {
id: MemoryId(Ulid::from(7u128)),
kind: MemoryKind::Fact,
body: b"source fact".to_vec(),
origin: None,
producer: None,
};
let mut projection = Projection::default();
projection.absorb_scoped(
&entry.encode().expect("entry encodes"),
Provenance::builder().conversation_id(conversation).build(),
StoredMemoryScope::default(),
Some(source.clone()),
);
let items = folded_memory_items(
&projection,
&MemoryScope::builder().conversation(conversation).build(),
&MemoryQuery::builder().build(),
);
assert_eq!(items[0].source, Some(source));
}
fn projection_with_feedback_on_the_older_item() -> (Projection, MemoryId, MemoryId) {
let conversation = ConversationId::new();
let prov = || Provenance::builder().conversation_id(conversation).build();
let mut projection = Projection::default();
let older = MemoryId::new();
let newer = MemoryId::new();
projection.absorb(&item_entry(older, b"older"), prov());
projection.absorb(&item_entry(newer, b"newer"), prov());
let feedback = MemoryLogEntry::Feedback {
target: older,
weight: 3.0,
conversation: None,
}
.encode()
.expect("a feedback entry encodes");
projection.absorb(&feedback, prov());
(projection, older, newer)
}
#[test]
fn given_feedback_when_folded_recall_is_recent_then_should_keep_recency_order() {
let (projection, older, newer) = projection_with_feedback_on_the_older_item();
let items = folded_memory_items(
&projection,
&MemoryScope::default(),
&MemoryQuery::builder()
.strategy(RecallStrategy::Recent)
.build(),
);
assert_eq!(
items.iter().map(|item| item.id).collect::<Vec<_>>(),
[newer, older]
);
assert!(items.iter().all(|item| item.score.is_none()));
assert!(items.iter().all(|item| item.signals.is_empty()));
}
#[test]
fn given_feedback_when_folded_recall_ranks_then_should_label_the_signal_auto() {
let (projection, older, _) = projection_with_feedback_on_the_older_item();
let items = folded_memory_items(
&projection,
&MemoryScope::default(),
&MemoryQuery::builder()
.strategy(RecallStrategy::Semantic)
.build(),
);
assert_eq!(items[0].id, older, "feedback promotes outside Recent");
assert_eq!(
items[0].signals,
[RecallSignal {
strategy: RecallStrategy::Auto,
rank: 0,
score: Some(3.0),
}]
);
}
#[tokio::test]
async fn given_feedback_when_vector_recall_is_recent_then_should_keep_recency_order() {
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::default();
let older = Memory::remember(&memory, &scope, b"cat".to_vec())
.await
.expect("remember older");
let newer = Memory::remember(&memory, &scope, b"dog".to_vec())
.await
.expect("remember newer");
Memory::improve(&memory, &scope, Feedback::new(older, 5.0))
.await
.expect("improve");
let items = Memory::recall(
&memory,
&scope,
&MemoryQuery::builder()
.strategy(RecallStrategy::Recent)
.build(),
)
.await
.expect("recent recall");
assert_eq!(
items.iter().map(|item| item.id).collect::<Vec<_>>(),
[newer, older]
);
assert!(items.iter().all(|item| item.signals.is_empty()));
}
#[tokio::test]
async fn given_a_builtin_memory_when_consolidated_then_should_append_a_durable_summary_kind() {
let memory = MemoryHandle::vector(WordEmbedder);
let conversation = ConversationId::new();
memory
.remember(b"cpu pressure".to_vec())
.scope(conversation)
.kind(MemoryKind::Message)
.send()
.await
.expect("message");
let scope = MemoryScope::builder().conversation(conversation).build();
let consolidator = DefaultConsolidator::new(&memory, 100).with_summarizer(StubSummarizer);
Consolidator::consolidate(&consolidator, &scope)
.await
.expect("consolidate");
let MemoryHandle::Vector(store) = &memory else {
panic!("vector")
};
let stored = store.items.lock().await;
let summary = stored
.iter()
.find(|entry| entry.item.kind == MemoryKind::Summary)
.expect("summary kind");
assert_eq!(summary.scope.lifetime, Lifetime::Durable);
assert_eq!(summary.item.payload, b"summary of 1");
}
#[tokio::test]
async fn given_recent_strategy_and_semantic_text_when_recalled_then_should_return_write_recency_without_embedding()
{
struct RecentEmbedder(std::sync::atomic::AtomicUsize);
impl Embedder for RecentEmbedder {
async fn embed(&self, _text: &str) -> Result<Vec<f32>, LaserError> {
self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(vec![1.0])
}
}
let memory = VectorMemory::new(RecentEmbedder(std::sync::atomic::AtomicUsize::new(0)));
let scope = MemoryScope::builder().stream("one".to_owned()).build();
Memory::remember(&memory, &scope, b"old".to_vec())
.await
.expect("old");
Memory::remember(&memory, &scope, b"new".to_vec())
.await
.expect("new");
Memory::remember(
&memory,
&MemoryScope::builder().stream("two".to_owned()).build(),
b"another stream".to_vec(),
)
.await
.expect("other");
let items = Memory::recall(
&memory,
&scope,
&MemoryQuery::builder()
.limit(1)
.semantic("ignored".to_owned())
.strategy(RecallStrategy::Recent)
.build(),
)
.await
.expect("recent");
assert_eq!(items[0].payload, b"new");
assert_eq!(
memory.embedder.0.load(std::sync::atomic::Ordering::SeqCst),
3
);
}
#[tokio::test]
async fn given_hybrid_memory_signals_when_ranked_then_should_report_additive_scores_and_each_signal_rank()
{
let memory = VectorMemory::new(WordEmbedder);
let scope = MemoryScope::default();
let first = Memory::remember(&memory, &scope, b"auth".to_vec())
.await
.expect("first");
let second = Memory::remember(&memory, &scope, b"auth latency".to_vec())
.await
.expect("second");
Memory::improve(&memory, &scope, Feedback::new(second, 1.5))
.await
.expect("feedback");
let items = Memory::recall(
&memory,
&scope,
&MemoryQuery::builder()
.semantic("auth".to_owned())
.strategy(RecallStrategy::Hybrid)
.build(),
)
.await
.expect("hybrid");
assert_eq!(items[0].id, second);
assert_eq!(items[1].id, first);
for item in &items {
assert_eq!(
item.score,
Some(
item.signals
.iter()
.map(|signal| signal.score.unwrap_or(0.0))
.sum()
)
);
assert!(
item.signals
.iter()
.any(|signal| signal.strategy == RecallStrategy::Semantic)
);
assert!(
item.signals
.iter()
.any(|signal| signal.strategy == RecallStrategy::Keyword)
);
}
let semantic = items[1]
.signals
.iter()
.find(|signal| signal.strategy == RecallStrategy::Semantic)
.expect("semantic");
assert_eq!(semantic.rank, 0);
assert_eq!(
items[0]
.signals
.iter()
.find(|signal| signal.strategy == RecallStrategy::Auto)
.expect("feedback")
.rank,
0
);
}
}