Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,8 @@
# MEMORY_MAX_MESSAGES_PER_RUN=200
# MEMORY_MIN_MESSAGES_PER_RUN=20
# MEMORY_MAX_QUEUED_RUNS=5000
# Visible-card lists the memory pipeline keeps in RAM (LRU); 0 disables it.
# MEMORY_CARD_CACHE_ENTRIES=4096
# MEMORY_MAX_INPUT_TOKENS=10000
# MEMORY_TOKENIZER_MODEL=google/gemma-4-26B-A4B-it
# MEMORY_TOKENIZER_FILE=
Expand Down
27 changes: 13 additions & 14 deletions crates/openplotva-app/src/dialog_jobs/input.rs
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,7 @@ pub struct PostgresDialogInputMaterializer {
history: PostgresHistoryStore,
members: PostgresChatMemberStore,
memory: Option<PostgresMemoryStore>,
memory_usage: Option<crate::memory_card_usage::MemoryCardUsageJournal>,
memory_embedder: Option<Arc<dyn EmbeddingProvider>>,
memory_embedding_dim: i32,
shield: Option<PostgresShieldStore>,
Expand Down Expand Up @@ -156,6 +157,7 @@ impl PostgresDialogInputMaterializer {
history,
members,
memory: None,
memory_usage: None,
memory_embedder: None,
memory_embedding_dim: 0,
shield: None,
Expand All @@ -168,9 +170,16 @@ impl PostgresDialogInputMaterializer {
}
}

/// Recall memory from `memory` and record recalled cards in `usage`, the
/// write-behind journal that applies them to Postgres.
#[must_use]
pub fn with_memory_store(mut self, memory: PostgresMemoryStore) -> Self {
pub fn with_memory_store(
mut self,
memory: PostgresMemoryStore,
usage: crate::memory_card_usage::MemoryCardUsageJournal,
) -> Self {
self.memory = Some(memory);
self.memory_usage = Some(usage);
self
}

Expand Down Expand Up @@ -757,19 +766,9 @@ impl PostgresDialogInputMaterializer {
.await
{
Ok(memory) => {
let card_ids = memory.cards.iter().map(|card| card.id).collect::<Vec<_>>();
if let Some(store) = self.memory.as_ref()
&& let Err(error) = store
.record_card_retrieval(&card_ids, OffsetDateTime::now_utc())
.await
{
tracing::warn!(
%error,
chat_id = params.chat_id,
user_id = params.user_id,
card_count = card_ids.len(),
"failed to record recalled memory cards"
);
if let Some(journal) = self.memory_usage.as_ref() {
let card_ids = memory.cards.iter().map(|card| card.id).collect::<Vec<_>>();
journal.record(&card_ids, OffsetDateTime::now_utc());
Comment thread
iamwavecut marked this conversation as resolved.
}
let captured = captured_memories_from_retrieval(&memory);
(
Expand Down
59 changes: 56 additions & 3 deletions crates/openplotva-app/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ mod lane_checks;
mod maintenance_api;
pub mod media;
pub mod members;
pub mod memory_card_usage;
pub mod memory_runtime;
pub mod message_gate;
pub mod model_routing;
Expand Down Expand Up @@ -350,6 +351,7 @@ struct RuntimeWorkers {
openrouter_free_gate: Option<Arc<openrouter_free_pool::OpenRouterFreeQuotaGate>>,
router_runtime: Option<Arc<model_routing::RouterRuntime>>,
dialog_worker_gauge: Option<Arc<dialog_workers::WorkerGauge>>,
memory_card_cache: Option<Arc<openplotva_storage::MemoryCardCache>>,
readiness_probes: Vec<ReadinessProbe>,
}

Expand Down Expand Up @@ -796,7 +798,10 @@ fn static_web_routes_from_config(
config.vip.chat_id,
)));
}
let memory_store = PostgresMemoryStore::new(clients.postgres.clone());
let mut memory_store = PostgresMemoryStore::new(clients.postgres.clone());
if let Some(cache) = &runtime_workers.memory_card_cache {
memory_store = memory_store.with_card_cache(Arc::clone(cache));
}
routes.memory_store = Some(memory_store.clone());
if config.memory.enabled
&& let (
Expand Down Expand Up @@ -10535,6 +10540,7 @@ async fn start_runtime_workers(
openrouter_free_gate: None,
router_runtime: None,
dialog_worker_gauge: None,
memory_card_cache: None,
readiness_probes: Vec::new(),
};
let readiness_general_postgres = service_clients.postgres.clone();
Expand Down Expand Up @@ -11237,7 +11243,52 @@ async fn start_runtime_workers(
.with_run_buffer(llm_run_buffer.clone()),
);
openplotva_llm::trace::set_observer(llm_observer);
let memory_store = PostgresMemoryStore::new(service_clients.postgres.clone());
let memory_card_cache = usize::try_from(config.memory.card_cache_entries)
.ok()
.filter(|entries| *entries > 0)
.map(|entries| {
Arc::new(openplotva_storage::MemoryCardCache::new(
entries,
openplotva_storage::MEMORY_CARD_CACHE_TTL,
))
});
workers.memory_card_cache = memory_card_cache.clone();
let memory_store = match &memory_card_cache {
Some(cache) => PostgresMemoryStore::new(service_clients.postgres.clone())
.with_card_cache(Arc::clone(cache)),
None => PostgresMemoryStore::new(service_clients.postgres.clone()),
};
let memory_card_usage = memory_card_usage::MemoryCardUsageJournal::new();
{
let journal = memory_card_usage.clone();
let store = memory_store.clone();
let cache = memory_card_cache.clone();
let stop = worker_stops.subscribe(RuntimeWorkerPhase::Ancillary);
let worker = tokio::spawn(async move {
let report = memory_card_usage::run_memory_card_usage_flush_worker_until(
journal,
store,
cache,
memory_card_usage::MEMORY_CARD_USAGE_FLUSH_INTERVAL,
wait_for_runtime_stop(stop),
)
.await;
tracing::info!(?report, "memory card usage flush worker stopped");
});
readiness_checks.push(ReadinessCheck::ok(
"memory_card_usage_flush",
format!(
"memory card uses flushed every {}s; visible-card cache {} entries",
memory_card_usage::MEMORY_CARD_USAGE_FLUSH_INTERVAL.as_secs(),
config.memory.card_cache_entries
),
));
worker_registry.register(
"memory-card-usage-flush",
RuntimeWorkerPhase::Ancillary,
worker,
);
}
let memory_restart_trigger = Arc::new(tokio::sync::Notify::new());
let virtual_dialog_store =
PostgresRuntimeVirtualDialogStore::new(service_clients.postgres.clone());
Expand Down Expand Up @@ -12944,7 +12995,8 @@ async fn start_runtime_workers(
),
);
if dialog_memory_context_enabled(config) {
dialog_materializer = dialog_materializer.with_memory_store(memory_store.clone());
dialog_materializer = dialog_materializer
.with_memory_store(memory_store.clone(), memory_card_usage.clone());
if let Some(embedder) = memory_query_embedder.clone() {
dialog_materializer = dialog_materializer
.with_memory_embedder(embedder, config.memory.embedding_dim);
Expand Down Expand Up @@ -14099,6 +14151,7 @@ async fn shutdown_runtime_workers(workers: RuntimeWorkers) {
openrouter_free_gate: _,
router_runtime: _,
dialog_worker_gauge: _,
memory_card_cache: _,
readiness_probes: _,
} = workers;
let started_at = Instant::now();
Expand Down
181 changes: 181 additions & 0 deletions crates/openplotva-app/src/memory_card_usage.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,181 @@
//! Write-behind journal of memory-card uses.
//!
//! Dialog turns record which cards they recalled without waiting for Postgres;
//! a worker applies the collected uses in one statement per interval and once
//! more on shutdown.

use std::{
collections::HashMap,
sync::{Arc, Mutex, MutexGuard},
time::Duration,
};

use openplotva_storage::{MemoryCardCache, MemoryCardUsage, PostgresMemoryStore};
use time::OffsetDateTime;

pub const MEMORY_CARD_USAGE_FLUSH_INTERVAL: Duration = Duration::from_secs(30);
/// Cache statistics are logged every this many flush ticks (5 minutes).
const CACHE_STATS_EVERY_TICKS: u64 = 10;

#[derive(Clone, Debug, Default)]
pub struct MemoryCardUsageJournal {
pending: Arc<Mutex<HashMap<i64, (OffsetDateTime, i64)>>>,
}

impl MemoryCardUsageJournal {
#[must_use]
pub fn new() -> Self {
Self::default()
}

/// Count one use of each card at `used_at`.
pub fn record(&self, card_ids: &[i64], used_at: OffsetDateTime) {
let mut pending = self.lock();
for card_id in card_ids.iter().copied().filter(|id| *id != 0) {
let entry = pending.entry(card_id).or_insert((used_at, 0));
entry.0 = entry.0.max(used_at);
entry.1 += 1;
}
}

fn drain(&self) -> Vec<MemoryCardUsage> {
self.lock()
.drain()
.map(|(card_id, (last_used_at, uses))| MemoryCardUsage {
card_id,
last_used_at,
uses,
})
.collect()
}

/// Put back uses whose flush failed, merging with uses recorded since.
fn restore(&self, usage: Vec<MemoryCardUsage>) {
let mut pending = self.lock();
for entry in usage {
let slot = pending
.entry(entry.card_id)
.or_insert((entry.last_used_at, 0));
slot.0 = slot.0.max(entry.last_used_at);
slot.1 += entry.uses;
}
}

fn lock(&self) -> MutexGuard<'_, HashMap<i64, (OffsetDateTime, i64)>> {
self.pending
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct MemoryCardUsageWorkerReport {
pub flushes: u64,
pub cards_updated: u64,
pub errors: u64,
}

/// Flush card uses every `interval` until `stop`, then flush what is left.
pub async fn run_memory_card_usage_flush_worker_until(
journal: MemoryCardUsageJournal,
store: PostgresMemoryStore,
card_cache: Option<Arc<MemoryCardCache>>,
interval: Duration,
stop: impl std::future::Future<Output = ()>,
) -> MemoryCardUsageWorkerReport {
let mut report = MemoryCardUsageWorkerReport::default();
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
ticker.tick().await;
tokio::pin!(stop);
let mut ticks = 0u64;
loop {
tokio::select! {
() = &mut stop => break,
_ = ticker.tick() => {
ticks += 1;
flush_card_usage(&journal, &store, &mut report).await;
if let Some(cache) = &card_cache
&& ticks.is_multiple_of(CACHE_STATS_EVERY_TICKS)
{
let stats = cache.stats();
tracing::info!(
entries = stats.entries,
bytes = stats.bytes,
hits = stats.hits,
misses = stats.misses,
expired = stats.expired,
evictions = stats.evictions,
invalidated = stats.invalidated,
"memory card cache stats"
);
}
}
}
}
flush_card_usage(&journal, &store, &mut report).await;
report
}

async fn flush_card_usage(
journal: &MemoryCardUsageJournal,
store: &PostgresMemoryStore,
report: &mut MemoryCardUsageWorkerReport,
) {
let usage = journal.drain();
if usage.is_empty() {
return;
}
match store.record_card_usage(&usage).await {
Ok(updated) => {
report.flushes += 1;
report.cards_updated += updated;
}
Err(error) => {
report.errors += 1;
tracing::warn!(%error, cards = usage.len(), "failed to flush memory card usage; keeping it for the next flush");
journal.restore(usage);
}
}
}

#[cfg(test)]
mod tests {
use super::*;

fn at(seconds: i64) -> OffsetDateTime {
OffsetDateTime::from_unix_timestamp(1_779_193_800 + seconds).expect("valid time")
}

fn by_card(usage: Vec<MemoryCardUsage>) -> Vec<(i64, OffsetDateTime, i64)> {
let mut usage = usage
.into_iter()
.map(|entry| (entry.card_id, entry.last_used_at, entry.uses))
.collect::<Vec<_>>();
usage.sort_unstable();
usage
}

#[test]
fn journal_counts_uses_and_keeps_the_latest_time_per_card() {
let journal = MemoryCardUsageJournal::new();
journal.record(&[1, 2, 0], at(10));
journal.record(&[2], at(5));
journal.record(&[1], at(20));

assert_eq!(by_card(journal.drain()), [(1, at(20), 2), (2, at(10), 2)]);
assert!(journal.drain().is_empty());
}

#[test]
fn failed_flush_is_merged_with_uses_recorded_meanwhile() {
let journal = MemoryCardUsageJournal::new();
journal.record(&[1], at(10));
let failed = journal.drain();
journal.record(&[1, 3], at(30));

journal.restore(failed);

assert_eq!(by_card(journal.drain()), [(1, at(30), 2), (3, at(30), 1)]);
}
}
Loading
Loading