halldyll_memory_model/storage/

redis.rs

//! Redis cache layer

use crate::core::{MemoryError, MemoryResult};
use redis::aio::ConnectionManager;
use redis::AsyncCommands;
use serde::{de::DeserializeOwned, Serialize};

/// Default TTL for cached items (1 hour)
const DEFAULT_TTL_SECONDS: u64 = 3600;

/// TTL for embeddings (24 hours)
const EMBEDDING_TTL_SECONDS: u64 = 86400;

/// TTL for session data (30 minutes)
const SESSION_TTL_SECONDS: u64 = 1800;

/// Redis caching operations
pub struct RedisCache {
    client: ConnectionManager,
    key_prefix: String,
}

impl RedisCache {
    /// Create new Redis cache
    pub fn new(client: ConnectionManager) -> Self {
        Self {
            client,
            key_prefix: "halldyll:".to_string(),
        }
    }

    /// Create new Redis cache with custom prefix
    pub fn with_prefix(client: ConnectionManager, prefix: impl Into<String>) -> Self {
        Self {
            client,
            key_prefix: prefix.into(),
        }
    }

    /// Build a full key with prefix
    fn build_key(&self, key: &str) -> String {
        format!("{}{}", self.key_prefix, key)
    }

    /// Cache an embedding vector
    pub async fn cache_embedding(&self, key: &str, embedding: &[f32]) -> MemoryResult<()> {
        let full_key = self.build_key(&format!("emb:{}", key));
        let json = serde_json::to_string(embedding)
            .map_err(|e| MemoryError::Serialization(e))?;

        let mut conn = self.client.clone();
        redis::cmd("SETEX")
            .arg(&full_key)
            .arg(EMBEDDING_TTL_SECONDS)
            .arg(&json)
            .query_async::<()>(&mut conn)
            .await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(())
    }

    /// Get cached embedding
    pub async fn get_embedding(&self, key: &str) -> MemoryResult<Option<Vec<f32>>> {
        let full_key = self.build_key(&format!("emb:{}", key));
        let mut conn = self.client.clone();

        let result: Option<String> = conn.get(&full_key).await
            .map_err(|e| MemoryError::Redis(e))?;

        match result {
            Some(json) => {
                let embedding: Vec<f32> = serde_json::from_str(&json)
                    .map_err(|e| MemoryError::Serialization(e))?;
                Ok(Some(embedding))
            }
            None => Ok(None),
        }
    }

    /// Cache a serializable value with default TTL
    pub async fn set<T: Serialize>(&self, key: &str, value: &T) -> MemoryResult<()> {
        self.set_with_ttl(key, value, DEFAULT_TTL_SECONDS).await
    }

    /// Cache a serializable value with custom TTL
    pub async fn set_with_ttl<T: Serialize>(
        &self,
        key: &str,
        value: &T,
        ttl_seconds: u64,
    ) -> MemoryResult<()> {
        let full_key = self.build_key(key);
        let json = serde_json::to_string(value)
            .map_err(|e| MemoryError::Serialization(e))?;

        let mut conn = self.client.clone();
        redis::cmd("SETEX")
            .arg(&full_key)
            .arg(ttl_seconds)
            .arg(&json)
            .query_async::<()>(&mut conn)
            .await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(())
    }

    /// Get a cached value
    pub async fn get<T: DeserializeOwned>(&self, key: &str) -> MemoryResult<Option<T>> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        let result: Option<String> = conn.get(&full_key).await
            .map_err(|e| MemoryError::Redis(e))?;

        match result {
            Some(json) => {
                let value: T = serde_json::from_str(&json)
                    .map_err(|e| MemoryError::Serialization(e))?;
                Ok(Some(value))
            }
            None => Ok(None),
        }
    }

    /// Invalidate cache key
    pub async fn invalidate(&self, key: &str) -> MemoryResult<()> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        conn.del::<_, ()>(&full_key).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(())
    }

    /// Invalidate all keys matching a pattern
    pub async fn invalidate_pattern(&self, pattern: &str) -> MemoryResult<u64> {
        let full_pattern = self.build_key(pattern);
        let mut conn = self.client.clone();

        // Get matching keys
        let keys: Vec<String> = redis::cmd("KEYS")
            .arg(&full_pattern)
            .query_async(&mut conn)
            .await
            .map_err(|e| MemoryError::Redis(e))?;

        if keys.is_empty() {
            return Ok(0);
        }

        // Delete all matching keys
        let count: u64 = conn.del(&keys).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(count)
    }

    /// Check if a key exists
    pub async fn exists(&self, key: &str) -> MemoryResult<bool> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        let exists: bool = conn.exists(&full_key).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(exists)
    }

    /// Set expiration on an existing key
    pub async fn expire(&self, key: &str, ttl_seconds: u64) -> MemoryResult<bool> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        let result: bool = conn.expire(&full_key, ttl_seconds as i64).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(result)
    }

    /// Get TTL of a key
    pub async fn ttl(&self, key: &str) -> MemoryResult<i64> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        let ttl: i64 = conn.ttl(&full_key).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(ttl)
    }

    /// Increment a counter
    pub async fn incr(&self, key: &str) -> MemoryResult<i64> {
        let full_key = self.build_key(key);
        let mut conn = self.client.clone();

        let value: i64 = conn.incr(&full_key, 1).await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(value)
    }

    /// Store session data
    pub async fn set_session<T: Serialize>(&self, session_id: &str, data: &T) -> MemoryResult<()> {
        let key = format!("session:{}", session_id);
        self.set_with_ttl(&key, data, SESSION_TTL_SECONDS).await
    }

    /// Get session data
    pub async fn get_session<T: DeserializeOwned>(&self, session_id: &str) -> MemoryResult<Option<T>> {
        let key = format!("session:{}", session_id);
        self.get(&key).await
    }

    /// Extend session TTL
    pub async fn extend_session(&self, session_id: &str) -> MemoryResult<bool> {
        let key = format!("session:{}", session_id);
        self.expire(&key, SESSION_TTL_SECONDS).await
    }

    /// Delete session
    pub async fn delete_session(&self, session_id: &str) -> MemoryResult<()> {
        let key = format!("session:{}", session_id);
        self.invalidate(&key).await
    }

    /// Cache user memory list
    pub async fn cache_user_memories(
        &self,
        user_id: &str,
        memory_type: &str,
        memory_ids: &[String],
    ) -> MemoryResult<()> {
        let key = format!("user:{}:memories:{}", user_id, memory_type);
        // Convert slice to Vec for serialization
        let ids_vec: Vec<String> = memory_ids.to_vec();
        self.set(&key, &ids_vec).await
    }

    /// Get cached user memory list
    pub async fn get_user_memories(
        &self,
        user_id: &str,
        memory_type: &str,
    ) -> MemoryResult<Option<Vec<String>>> {
        let key = format!("user:{}:memories:{}", user_id, memory_type);
        self.get(&key).await
    }

    /// Invalidate user memory cache
    pub async fn invalidate_user_cache(&self, user_id: &str) -> MemoryResult<u64> {
        let pattern = format!("user:{}:*", user_id);
        self.invalidate_pattern(&pattern).await
    }

    /// Ping Redis to check connectivity
    pub async fn ping(&self) -> MemoryResult<bool> {
        let mut conn = self.client.clone();

        let result: String = redis::cmd("PING")
            .query_async(&mut conn)
            .await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(result == "PONG")
    }

    /// Get Redis info
    pub async fn info(&self) -> MemoryResult<String> {
        let mut conn = self.client.clone();

        let info: String = redis::cmd("INFO")
            .query_async(&mut conn)
            .await
            .map_err(|e| MemoryError::Redis(e))?;

        Ok(info)
    }
}

#[cfg(test)]
mod tests {
    // Note: These tests require a running Redis instance
    // They are marked as ignored by default

    #[test]
    fn test_build_key() {
        // Mock test without actual Redis connection
        let key_prefix = "halldyll:";
        let key = "test";
        let full_key = format!("{}{}", key_prefix, key);
        assert_eq!(full_key, "halldyll:test");
    }

    #[test]
    fn test_embedding_key_format() {
        let key_prefix = "halldyll:";
        let key = "my_text";
        let full_key = format!("{}emb:{}", key_prefix, key);
        assert_eq!(full_key, "halldyll:emb:my_text");
    }
}