use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::chunk::NewMemoryChunk;
use crate::error::Result;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum WriteDecision {
Allow,
Redact {
reason: String,
},
}
#[async_trait]
pub trait MemoryWriteHook: Send + Sync {
async fn before_write(&self, _chunk: &NewMemoryChunk) -> Result<WriteDecision> {
Ok(WriteDecision::Allow)
}
}
pub struct ClosureHook<F>(pub F)
where
F: Fn(&NewMemoryChunk) -> WriteDecision + Send + Sync + 'static;
#[async_trait]
impl<F> MemoryWriteHook for ClosureHook<F>
where
F: Fn(&NewMemoryChunk) -> WriteDecision + Send + Sync + 'static,
{
async fn before_write(&self, chunk: &NewMemoryChunk) -> Result<WriteDecision> {
Ok((self.0)(chunk))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chunk::{ChunkKind, ChunkSource};
fn nc(content: &str) -> NewMemoryChunk {
NewMemoryChunk {
kind: ChunkKind::Chat,
source: ChunkSource::Embedded,
session_id: None,
project_root: None,
caller_id: "test".into(),
content: content.into(),
metadata: serde_json::Value::Null,
importance: None,
shareable: false,
pinned: false,
}
}
#[tokio::test]
async fn closure_hook_redacts_matching_content() {
let hook = ClosureHook(|c: &NewMemoryChunk| {
if c.content.contains("secret") {
WriteDecision::Redact {
reason: "secret keyword".into(),
}
} else {
WriteDecision::Allow
}
});
let chunk = nc("a totally innocent chat");
assert_eq!(
hook.before_write(&chunk).await.unwrap(),
WriteDecision::Allow
);
let chunk = nc("contains a secret");
match hook.before_write(&chunk).await.unwrap() {
WriteDecision::Redact { reason } => assert_eq!(reason, "secret keyword"),
other => panic!("expected Redact, got {other:?}"),
}
}
#[tokio::test]
async fn default_hook_allows_everything() {
struct Default;
#[async_trait]
impl MemoryWriteHook for Default {}
let chunk = nc("anything");
assert_eq!(
Default.before_write(&chunk).await.unwrap(),
WriteDecision::Allow
);
}
}