agentlib-memory 0.1.0

use agentlib_core::{
    MemoryProvider, MemoryReadOptions, MemoryWriteOptions, ModelMessage, async_trait,
};
use anyhow::Result;
use std::collections::HashSet;

pub enum ReadStrategy {
    FirstHit,
    Merge,
}

pub struct CompositeMemory {
    providers: Vec<Box<dyn MemoryProvider>>,
    read_strategy: ReadStrategy,
}

impl CompositeMemory {
    pub fn new(providers: Vec<Box<dyn MemoryProvider>>, read_strategy: ReadStrategy) -> Self {
        Self {
            providers,
            read_strategy,
        }
    }
}

#[async_trait]
impl MemoryProvider for CompositeMemory {
    async fn read(&self, options: MemoryReadOptions) -> Result<Vec<ModelMessage>> {
        match self.read_strategy {
            ReadStrategy::FirstHit => {
                for provider in &self.providers {
                    let messages = provider.read(options.clone()).await?;
                    if !messages.is_empty() {
                        return Ok(messages);
                    }
                }
                Ok(Vec::new())
            }
            ReadStrategy::Merge => {
                let mut merged = Vec::new();
                let mut seen = HashSet::new();

                for provider in &self.providers {
                    let messages = provider.read(options.clone()).await?;
                    for msg in messages {
                        let key = format!(
                            "{:?}:{}",
                            msg.role,
                            msg.content.chars().take(100).collect::<String>()
                        );
                        if !seen.contains(&key) {
                            seen.insert(key);
                            merged.push(msg);
                        }
                    }
                }
                Ok(merged)
            }
        }
    }

    async fn write(&self, messages: Vec<ModelMessage>, options: MemoryWriteOptions) -> Result<()> {
        for provider in &self.providers {
            provider.write(messages.clone(), options.clone()).await?;
        }
        Ok(())
    }

    async fn clear(&self, session_id: Option<&str>) -> Result<()> {
        for provider in &self.providers {
            provider.clear(session_id).await?;
        }
        Ok(())
    }
}