use std::sync::Arc;
use async_trait::async_trait;
use uuid::Uuid;
use mr_common::SimilarMemory;
use crate::llm::{LlmClient, LlmMessage};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DedupAction {
Add,
Skip,
Update(Uuid),
}
#[derive(Debug, Clone)]
pub struct DedupDecision {
pub action: DedupAction,
pub candidates: Vec<SimilarMemory>,
}
impl DedupDecision {
pub fn add() -> Self {
Self {
action: DedupAction::Add,
candidates: Vec::new(),
}
}
}
#[async_trait]
pub trait DedupDecider: Send + Sync {
async fn decide(&self, content: &str, candidates: &[SimilarMemory]) -> DedupDecision;
}
#[derive(Debug, Clone)]
pub struct RuleDedupDecider {
duplicate_threshold: f32,
update_threshold: f32,
}
impl RuleDedupDecider {
pub fn new(duplicate_threshold: f32, update_threshold: f32) -> Self {
Self {
duplicate_threshold: duplicate_threshold.clamp(0.0, 1.0),
update_threshold: update_threshold.clamp(0.0, 1.0),
}
}
pub fn from_config(config: &mr_common::DedupConfig) -> Self {
Self::new(config.duplicate_threshold, config.update_threshold)
}
}
#[async_trait]
impl DedupDecider for RuleDedupDecider {
async fn decide(&self, _content: &str, candidates: &[SimilarMemory]) -> DedupDecision {
if let Some(best) = candidates.first() {
if best.score >= self.duplicate_threshold {
return DedupDecision {
action: DedupAction::Skip,
candidates: candidates.to_vec(),
};
}
if best.score >= self.update_threshold {
return DedupDecision {
action: DedupAction::Update(best.memory_id),
candidates: candidates.to_vec(),
};
}
}
DedupDecision {
action: DedupAction::Add,
candidates: candidates.to_vec(),
}
}
}
#[allow(dead_code)]
pub struct LlmDedupDecider {
llm: Arc<dyn LlmClient>,
rule: RuleDedupDecider,
}
impl LlmDedupDecider {
pub fn new(llm: Arc<dyn LlmClient>, duplicate_threshold: f32, update_threshold: f32) -> Self {
Self {
llm,
rule: RuleDedupDecider::new(duplicate_threshold, update_threshold),
}
}
fn build_prompt(&self, content: &str, candidates: &[SimilarMemory]) -> String {
let mut prompt = String::from(
"你是记忆库维护助手。判断新记忆与已有记忆的关系,输出 JSON:\
{\"action\": \"add|update|delete|noop\", \"target_id\": \"...\"}\n\n",
);
prompt.push_str("新记忆:\n");
prompt.push_str(content);
prompt.push('\n');
prompt.push_str("已有记忆:\n");
for c in candidates {
prompt.push_str(&format!(
"- [{:?}] {} (score={:.2})\n",
c.memory_id, c.content, c.score
));
}
prompt
}
}
#[async_trait]
impl DedupDecider for LlmDedupDecider {
async fn decide(&self, content: &str, candidates: &[SimilarMemory]) -> DedupDecision {
let _ = (self
.llm
.chat(&[LlmMessage::user(self.build_prompt(content, candidates))])
.await,);
self.rule.decide(content, candidates).await
}
}
#[cfg(test)]
mod tests {
use super::*;
fn candidate(id: u32, score: f32) -> SimilarMemory {
SimilarMemory {
memory_id: Uuid::from_u128(id as u128),
content: format!("memory {}", id),
score,
}
}
#[tokio::test]
async fn test_duplicate_skip() {
let decider = RuleDedupDecider::new(0.92, 0.85);
let decision = decider.decide("new content", &[candidate(1, 0.95)]).await;
assert_eq!(decision.action, DedupAction::Skip);
}
#[tokio::test]
async fn test_update_merge() {
let decider = RuleDedupDecider::new(0.92, 0.85);
let decision = decider.decide("new content", &[candidate(1, 0.88)]).await;
assert_eq!(decision.action, DedupAction::Update(Uuid::from_u128(1)));
}
#[tokio::test]
async fn test_add_when_below_update_threshold() {
let decider = RuleDedupDecider::new(0.92, 0.85);
let decision = decider.decide("new content", &[candidate(1, 0.80)]).await;
assert_eq!(decision.action, DedupAction::Add);
}
#[tokio::test]
async fn test_add_when_no_candidates() {
let decider = RuleDedupDecider::new(0.92, 0.85);
let decision = decider.decide("new content", &[]).await;
assert_eq!(decision.action, DedupAction::Add);
assert!(decision.candidates.is_empty());
}
#[tokio::test]
async fn test_candidates_kept_in_decision() {
let decider = RuleDedupDecider::new(0.92, 0.85);
let candidates = vec![candidate(1, 0.95), candidate(2, 0.90)];
let decision = decider.decide("new content", &candidates).await;
assert_eq!(decision.candidates.len(), 2);
assert_eq!(decision.candidates[0].memory_id, Uuid::from_u128(1));
}
#[tokio::test]
async fn test_threshold_clamped() {
let decider = RuleDedupDecider::new(2.0, -1.0);
assert_eq!(decider.duplicate_threshold, 1.0);
assert_eq!(decider.update_threshold, 0.0);
let decision = decider.decide("x", &[candidate(1, 0.01)]).await;
assert_eq!(decision.action, DedupAction::Update(Uuid::from_u128(1)));
}
#[tokio::test]
async fn test_llm_decider_falls_back_to_rule() {
let llm: Arc<dyn LlmClient> = Arc::new(crate::llm::MockLlmClient::new("ignored"));
let decider = LlmDedupDecider::new(llm, 0.92, 0.85);
let decision = decider.decide("new content", &[candidate(1, 0.95)]).await;
assert_eq!(decision.action, DedupAction::Skip);
}
}