use std::sync::Arc;
use rig::completion::Message;
use rig::memory::{Compactor, MemoryError};
use rig::wasm_compat::WasmBoxedFuture;
use crate::dedup::DedupSet;
use crate::frame_writer::{commit_if_each_turn, write_frame};
use crate::hook::{MemoryConfig, render_message_text};
use crate::store::MemvidStore;
pub struct MemvidStoringCompactor<C> {
store: MemvidStore,
inner: C,
config: MemoryConfig,
dedup: Arc<DedupSet>,
}
impl<C: Clone> Clone for MemvidStoringCompactor<C> {
fn clone(&self) -> Self {
Self {
store: self.store.clone(),
inner: self.inner.clone(),
config: self.config.clone(),
dedup: Arc::clone(&self.dedup),
}
}
}
impl<C: std::fmt::Debug> std::fmt::Debug for MemvidStoringCompactor<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MemvidStoringCompactor")
.field("inner", &self.inner)
.field("config", &self.config)
.field("dedup", &self.dedup)
.finish_non_exhaustive()
}
}
impl<C> MemvidStoringCompactor<C> {
pub fn new(store: MemvidStore, inner: C, config: MemoryConfig) -> Self {
Self {
store,
inner,
config,
dedup: Arc::new(DedupSet::new()),
}
}
pub fn with_defaults(store: MemvidStore, inner: C) -> Self {
Self::new(store, inner, MemoryConfig::default())
}
pub fn inner(&self) -> &C {
&self.inner
}
pub fn store(&self) -> &MemvidStore {
&self.store
}
pub fn config(&self) -> &MemoryConfig {
&self.config
}
pub fn dedup_snapshot(&self) -> Result<Vec<String>, MemoryError> {
self.dedup
.snapshot()
.map_err(|err| MemoryError::backend(Box::new(err)))
}
pub fn load_dedup_snapshot(&self, hexes: &[String]) -> Result<(), MemoryError> {
self.dedup
.extend_from_snapshot(hexes)
.map_err(|err| MemoryError::backend(Box::new(err)))
}
}
impl<C> Compactor for MemvidStoringCompactor<C>
where
C: Compactor,
{
type Artifact = C::Artifact;
fn compact<'a>(
&'a self,
conversation_id: &'a str,
evicted: &'a [Message],
carry_over: Option<&'a Self::Artifact>,
) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
Box::pin(async move {
let artifact = self
.inner
.compact(conversation_id, evicted, carry_over)
.await?;
if let Some(rendered) = render_message_text(&artifact.clone().into()) {
let written = write_frame(
&self.store,
&self.config,
&self.dedup,
crate::metadata::FrameKind::CompactionSummary,
conversation_id,
"system",
&rendered,
)
.await?;
if written {
commit_if_each_turn(&self.store, &self.config).await?;
}
#[cfg(feature = "observe")]
{
let evicted_bytes: usize = evicted
.iter()
.filter_map(render_message_text)
.map(|s| s.len())
.sum();
rig_tap::emit_kind(
conversation_id,
rig_tap::EventKind::ContextCompacted {
evicted_count: evicted.len(),
evicted_bytes,
carry_over: carry_over.is_some(),
summary_bytes: rendered.len(),
},
);
}
}
Ok(artifact)
})
}
}