ambi_memory/provider/
kv.rs

1//! Exact key-value memory provider for reflexion diaries, settings, and state flags.
2
3use crate::error::Result;
4use async_trait::async_trait;
5use std::collections::HashMap;
6use std::sync::Arc;
7use tokio::sync::RwLock;
8
9/// Persistent key-value memory interface.
10///
11/// Each value is scoped to a `session_id`, allowing independent memory
12/// namespaces per conversation.
13#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
14#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
15pub trait KvMemoryProvider: Send + Sync {
16    /// Store a value under the given key for a session.
17    async fn store(&self, session_id: &str, key: &str, value: &str) -> Result<()>;
18    /// Retrieve a single value by key.
19    async fn retrieve(&self, session_id: &str, key: &str) -> Result<Option<String>>;
20    /// Retrieve all key-value pairs for a session.
21    async fn retrieve_all(&self, session_id: &str) -> Result<HashMap<String, String>>;
22    /// Remove all data associated with a session.
23    async fn clear_session(&self, session_id: &str) -> Result<()>;
24}
25
26/// In-memory implementation of `KvMemoryProvider` (for testing and single-node deployments).
27#[derive(Clone, Default)]
28pub struct InMemoryKvProvider {
29    store: Arc<RwLock<HashMap<String, HashMap<String, String>>>>,
30}
31
32impl InMemoryKvProvider {
33    /// Creates an empty in-memory KV store.
34    pub fn new() -> Self {
35        Self {
36            store: Arc::new(RwLock::new(HashMap::new())),
37        }
38    }
39}
40
41#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
42#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
43impl KvMemoryProvider for InMemoryKvProvider {
44    async fn store(&self, session_id: &str, key: &str, value: &str) -> Result<()> {
45        let mut lock = self.store.write().await;
46        lock.entry(session_id.to_string())
47            .or_default()
48            .insert(key.to_string(), value.to_string());
49        Ok(())
50    }
51    async fn retrieve(&self, session_id: &str, key: &str) -> Result<Option<String>> {
52        Ok(self
53            .store
54            .read()
55            .await
56            .get(session_id)
57            .and_then(|s| s.get(key).cloned()))
58    }
59    async fn retrieve_all(&self, session_id: &str) -> Result<HashMap<String, String>> {
60        Ok(self
61            .store
62            .read()
63            .await
64            .get(session_id)
65            .cloned()
66            .unwrap_or_default())
67    }
68    async fn clear_session(&self, session_id: &str) -> Result<()> {
69        self.store.write().await.remove(session_id);
70        Ok(())
71    }
72}
ambi_memory/provider/kv.rs

ambi_memory/provider/
kv.rs
