enki-runtime 0.1.4

//! Memory backend configuration from TOML files.

use serde::{Deserialize, Serialize};

/// Memory backend type.
///
/// Specifies which memory backend implementation to use.
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum MemoryBackendType {
    /// In-memory storage (fast, no persistence)
    #[default]
    InMemory,
    /// SQLite-based storage (persistent, requires 'sqlite' feature)
    Sqlite,
    /// Redis-based storage (distributed, requires 'redis' feature)
    Redis,
    /// Qdrant vector database (requires 'qdrant' feature)
    Qdrant,
    /// ChromaDB vector database (requires 'chromadb' feature)
    #[serde(alias = "chroma")]
    ChromaDB,
    /// Pinecone vector database (requires 'pinecone' feature)
    Pinecone,
}

/// Configuration for memory backends.
///
/// This struct can be used at the mesh level (shared by all agents) or
/// at the agent level (per-agent configuration).
///
/// # Examples
///
/// ## InMemory (default)
/// ```toml
/// [memory]
/// backend = "inmemory"
/// max_entries = 500
/// ttl_seconds = 3600
/// ```
///
/// ## SQLite (persistent)
/// ```toml
/// [memory]
/// backend = "sqlite"
/// path = "./data/knowledge.db"
/// ```
///
/// ## Redis (distributed)
/// ```toml
/// [memory]
/// backend = "redis"
/// url = "redis://localhost:6379"
/// prefix = "myapp"
/// ttl_seconds = 7200
/// ```
///
/// ## Qdrant (vector database)
/// ```toml
/// [memory]
/// backend = "qdrant"
/// grpc_url = "http://localhost:6334"
/// collection = "agent_memory"
/// dimensions = 384
/// embedding_model = "ollama::mxbai-embed-large"
/// ```
///
/// ## ChromaDB (vector database)
/// ```toml
/// [memory]
/// backend = "chromadb"
/// url = "http://localhost:8000"
/// collection = "agent_memory"
/// use_server_embeddings = true
/// ```
///
/// ## Pinecone (cloud vector database)
/// ```toml
/// [memory]
/// backend = "pinecone"
/// api_key = "${PINECONE_API_KEY}"
/// collection = "agent-memory"
/// dimensions = 1536
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MemoryConfig {
    /// Backend type: "inmemory", "sqlite", "redis", "qdrant", "chromadb", or "pinecone"
    #[serde(default)]
    pub backend: MemoryBackendType,

    /// Path for SQLite database file.
    /// Use ":memory:" for in-memory SQLite.
    /// Only used when backend = "sqlite".
    #[serde(default)]
    pub path: Option<String>,

    /// Connection URL for Redis or ChromaDB.
    /// Example: "redis://localhost:6379" or "http://localhost:8000"
    /// Used when backend = "redis" or "chromadb".
    #[serde(default)]
    pub url: Option<String>,

    /// GRPC URL for Qdrant.
    /// Example: "http://localhost:6334"
    /// Only used when backend = "qdrant".
    #[serde(default)]
    pub grpc_url: Option<String>,

    /// Collection/index name for vector databases.
    /// Used by Qdrant, ChromaDB, and Pinecone.
    #[serde(default)]
    pub collection: Option<String>,

    /// Vector dimensions for embeddings.
    /// Required for Qdrant and Pinecone when creating new collections.
    #[serde(default)]
    pub dimensions: Option<usize>,

    /// API key for cloud services (Pinecone, etc.).
    /// Can use environment variable syntax: "${PINECONE_API_KEY}"
    #[serde(default)]
    pub api_key: Option<String>,

    /// Embedding model identifier (e.g., "ollama::mxbai-embed-large").
    /// Used for generating embeddings if not using server-side embeddings.
    #[serde(default)]
    pub embedding_model: Option<String>,

    /// Whether to use the vector DB's server-side embedding feature.
    /// ChromaDB supports this with configured embedding functions.
    #[serde(default)]
    pub use_server_embeddings: Option<bool>,

    /// Key prefix for Redis namespace isolation.
    /// Only used when backend = "redis".
    #[serde(default)]
    pub prefix: Option<String>,

    /// Maximum number of entries for InMemory backend.
    /// When exceeded, oldest entries are evicted.
    /// Only used when backend = "inmemory".
    #[serde(default)]
    pub max_entries: Option<usize>,

    /// Default TTL (time-to-live) in seconds for entries.
    /// Entries will be automatically removed after this time.
    /// Supported by all backends.
    #[serde(default)]
    pub ttl_seconds: Option<i64>,
}

impl MemoryConfig {
    /// Create a new default memory config (InMemory backend).
    pub fn new() -> Self {
        Self::default()
    }

    /// Create an InMemory backend configuration.
    pub fn in_memory() -> Self {
        Self {
            backend: MemoryBackendType::InMemory,
            ..Default::default()
        }
    }

    /// Create a SQLite backend configuration.
    pub fn sqlite(path: impl Into<String>) -> Self {
        Self {
            backend: MemoryBackendType::Sqlite,
            path: Some(path.into()),
            ..Default::default()
        }
    }

    /// Create a Redis backend configuration.
    pub fn redis(url: impl Into<String>) -> Self {
        Self {
            backend: MemoryBackendType::Redis,
            url: Some(url.into()),
            ..Default::default()
        }
    }

    /// Create a Qdrant backend configuration.
    ///
    /// # Arguments
    /// * `grpc_url` - Qdrant gRPC URL (e.g., "http://localhost:6334")
    /// * `collection` - Collection name
    /// * `dimensions` - Vector dimensions
    pub fn qdrant(
        grpc_url: impl Into<String>,
        collection: impl Into<String>,
        dimensions: usize,
    ) -> Self {
        Self {
            backend: MemoryBackendType::Qdrant,
            grpc_url: Some(grpc_url.into()),
            collection: Some(collection.into()),
            dimensions: Some(dimensions),
            ..Default::default()
        }
    }

    /// Create a ChromaDB backend configuration.
    ///
    /// # Arguments
    /// * `url` - ChromaDB URL (e.g., "http://localhost:8000")
    /// * `collection` - Collection name
    pub fn chromadb(url: impl Into<String>, collection: impl Into<String>) -> Self {
        Self {
            backend: MemoryBackendType::ChromaDB,
            url: Some(url.into()),
            collection: Some(collection.into()),
            ..Default::default()
        }
    }

    /// Create a Pinecone backend configuration.
    ///
    /// # Arguments
    /// * `api_key` - Pinecone API key
    /// * `collection` - Index name
    /// * `dimensions` - Vector dimensions
    pub fn pinecone(
        api_key: impl Into<String>,
        collection: impl Into<String>,
        dimensions: usize,
    ) -> Self {
        Self {
            backend: MemoryBackendType::Pinecone,
            api_key: Some(api_key.into()),
            collection: Some(collection.into()),
            dimensions: Some(dimensions),
            ..Default::default()
        }
    }

    /// Set max entries limit.
    pub fn with_max_entries(mut self, max: usize) -> Self {
        self.max_entries = Some(max);
        self
    }

    /// Set default TTL in seconds.
    pub fn with_ttl_seconds(mut self, seconds: i64) -> Self {
        self.ttl_seconds = Some(seconds);
        self
    }

    /// Set Redis key prefix.
    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.prefix = Some(prefix.into());
        self
    }

    /// Set collection/index name.
    pub fn with_collection(mut self, collection: impl Into<String>) -> Self {
        self.collection = Some(collection.into());
        self
    }

    /// Set vector dimensions.
    pub fn with_dimensions(mut self, dimensions: usize) -> Self {
        self.dimensions = Some(dimensions);
        self
    }

    /// Set embedding model.
    pub fn with_embedding_model(mut self, model: impl Into<String>) -> Self {
        self.embedding_model = Some(model.into());
        self
    }

    /// Enable or disable server-side embeddings.
    pub fn with_server_embeddings(mut self, enabled: bool) -> Self {
        self.use_server_embeddings = Some(enabled);
        self
    }

    /// Set API key for cloud services.
    pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
        self.api_key = Some(api_key.into());
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_inmemory_config() {
        let toml = r#"
            backend = "inmemory"
            max_entries = 500
            ttl_seconds = 3600
        "#;

        let config: MemoryConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.backend, MemoryBackendType::InMemory);
        assert_eq!(config.max_entries, Some(500));
        assert_eq!(config.ttl_seconds, Some(3600));
    }

    #[test]
    fn test_parse_sqlite_config() {
        let toml = r#"
            backend = "sqlite"
            path = "./data/knowledge.db"
        "#;

        let config: MemoryConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.backend, MemoryBackendType::Sqlite);
        assert_eq!(config.path, Some("./data/knowledge.db".to_string()));
    }

    #[test]
    fn test_parse_redis_config() {
        let toml = r#"
            backend = "redis"
            url = "redis://localhost:6379"
            prefix = "myapp"
            ttl_seconds = 7200
        "#;

        let config: MemoryConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.backend, MemoryBackendType::Redis);
        assert_eq!(config.url, Some("redis://localhost:6379".to_string()));
        assert_eq!(config.prefix, Some("myapp".to_string()));
        assert_eq!(config.ttl_seconds, Some(7200));
    }

    #[test]
    fn test_default_is_inmemory() {
        let toml = "";
        let config: MemoryConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.backend, MemoryBackendType::InMemory);
    }

    #[test]
    fn test_builder_methods() {
        let config = MemoryConfig::in_memory()
            .with_max_entries(1000)
            .with_ttl_seconds(3600);

        assert_eq!(config.backend, MemoryBackendType::InMemory);
        assert_eq!(config.max_entries, Some(1000));
        assert_eq!(config.ttl_seconds, Some(3600));

        let config = MemoryConfig::redis("redis://localhost:6379")
            .with_prefix("test")
            .with_ttl_seconds(7200);

        assert_eq!(config.backend, MemoryBackendType::Redis);
        assert_eq!(config.url, Some("redis://localhost:6379".to_string()));
        assert_eq!(config.prefix, Some("test".to_string()));
    }
}