//! OpenViking Memory Storage - Native Rust Implementation //! //! Provides native Rust memory storage using SqliteStorage with TF-IDF semantic search. //! This is a self-contained implementation that doesn't require external Python or CLI dependencies. //! //! Features: //! - SQLite persistence with FTS5 full-text search //! - TF-IDF semantic scoring //! - Token budget control //! - Automatic memory indexing use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::OnceCell; use zclaw_growth::{ FindOptions, MemoryEntry, MemoryType, PromptInjector, RetrievalResult, SqliteStorage, VikingStorage, }; // === Types === #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct VikingStatus { pub available: bool, pub version: Option, pub data_dir: Option, pub error: Option, } #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct VikingResource { pub uri: String, pub name: String, #[serde(rename = "type")] pub resource_type: String, pub size: Option, pub modified_at: Option, } #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct VikingFindResult { pub uri: String, pub score: f64, pub content: String, pub level: String, pub overview: Option, } #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct VikingGrepResult { pub uri: String, pub line: u32, pub content: String, pub match_start: u32, pub match_end: u32, } #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct VikingAddResult { pub uri: String, pub status: String, } // === Global Storage Instance === /// Global storage instance static STORAGE: OnceCell> = OnceCell::const_new(); /// Get the storage directory path fn get_storage_dir() -> PathBuf { // Use platform-specific data directory if let Some(data_dir) = dirs::data_dir() { data_dir.join("zclaw").join("memories") } else { // Fallback to current directory PathBuf::from("./zclaw_data/memories") } } /// Initialize the storage (should be called once at startup) pub async fn init_storage() -> Result<(), String> { let storage_dir = get_storage_dir(); let db_path = storage_dir.join("memories.db"); tracing::info!("[VikingCommands] Initializing storage at {:?}", db_path); let storage = SqliteStorage::new(&db_path) .await .map_err(|e| format!("Failed to initialize storage: {}", e))?; let _ = STORAGE.set(Arc::new(storage)); tracing::info!("[VikingCommands] Storage initialized successfully"); Ok(()) } /// Get the storage instance (public for use by other modules) pub async fn get_storage() -> Result, String> { STORAGE .get() .cloned() .ok_or_else(|| "Storage not initialized. Call init_storage() first.".to_string()) } /// Get storage directory for status fn get_data_dir_string() -> Option { get_storage_dir().to_str().map(|s| s.to_string()) } // === Tauri Commands === /// Check if memory storage is available #[tauri::command] pub async fn viking_status() -> Result { match get_storage().await { Ok(storage) => { // Try a simple query to verify storage is working let _ = storage .find("", FindOptions::default()) .await .map_err(|e| format!("Storage health check failed: {}", e))?; Ok(VikingStatus { available: true, version: Some("0.1.0-native".to_string()), data_dir: get_data_dir_string(), error: None, }) } Err(e) => Ok(VikingStatus { available: false, version: None, data_dir: get_data_dir_string(), error: Some(e), }), } } /// Add a memory entry #[tauri::command] pub async fn viking_add(uri: String, content: String) -> Result { let storage = get_storage().await?; // Parse URI to extract agent_id, memory_type, and category // Expected format: agent://{agent_id}/{type}/{category} let (agent_id, memory_type, category) = parse_uri(&uri)?; let entry = MemoryEntry::new(&agent_id, memory_type, &category, content); storage .store(&entry) .await .map_err(|e| format!("Failed to store memory: {}", e))?; Ok(VikingAddResult { uri, status: "added".to_string(), }) } /// Add a memory with metadata #[tauri::command] pub async fn viking_add_with_metadata( uri: String, content: String, keywords: Vec, importance: Option, ) -> Result { let storage = get_storage().await?; let (agent_id, memory_type, category) = parse_uri(&uri)?; let mut entry = MemoryEntry::new(&agent_id, memory_type, &category, content); entry.keywords = keywords; if let Some(imp) = importance { entry.importance = imp.min(10).max(1); } storage .store(&entry) .await .map_err(|e| format!("Failed to store memory: {}", e))?; Ok(VikingAddResult { uri, status: "added".to_string(), }) } /// Find memories by semantic search #[tauri::command] pub async fn viking_find( query: String, scope: Option, limit: Option, ) -> Result, String> { let storage = get_storage().await?; let options = FindOptions { scope, limit, min_similarity: Some(0.1), }; let entries = storage .find(&query, options) .await .map_err(|e| format!("Failed to search memories: {}", e))?; Ok(entries .into_iter() .enumerate() .map(|(i, entry)| VikingFindResult { uri: entry.uri, score: 1.0 - (i as f64 * 0.1), // Simple scoring based on rank content: entry.content, level: "L1".to_string(), overview: None, }) .collect()) } /// Grep memories by pattern (uses FTS5) #[tauri::command] pub async fn viking_grep( pattern: String, uri: Option, _case_sensitive: Option, limit: Option, ) -> Result, String> { let storage = get_storage().await?; let scope = uri.as_ref().and_then(|u| { // Extract agent scope from URI u.strip_prefix("agent://") .and_then(|s| s.split('/').next()) .map(|agent| format!("agent://{}", agent)) }); let options = FindOptions { scope, limit, min_similarity: Some(0.05), // Lower threshold for grep }; let entries = storage .find(&pattern, options) .await .map_err(|e| format!("Failed to grep memories: {}", e))?; Ok(entries .into_iter() .flat_map(|entry| { // Find matching lines entry .content .lines() .enumerate() .filter(|(_, line)| { line.to_lowercase() .contains(&pattern.to_lowercase()) }) .map(|(i, line)| VikingGrepResult { uri: entry.uri.clone(), line: (i + 1) as u32, content: line.to_string(), match_start: line.find(&pattern).unwrap_or(0) as u32, match_end: (line.find(&pattern).unwrap_or(0) + pattern.len()) as u32, }) .collect::>() }) .take(limit.unwrap_or(100)) .collect()) } /// List memories at a path #[tauri::command] pub async fn viking_ls(path: String) -> Result, String> { let storage = get_storage().await?; let entries = storage .find_by_prefix(&path) .await .map_err(|e| format!("Failed to list memories: {}", e))?; Ok(entries .into_iter() .map(|entry| VikingResource { uri: entry.uri.clone(), name: entry .uri .rsplit('/') .next() .unwrap_or(&entry.uri) .to_string(), resource_type: entry.memory_type.to_string(), size: Some(entry.content.len() as u64), modified_at: Some(entry.last_accessed.to_rfc3339()), }) .collect()) } /// Read memory content #[tauri::command] pub async fn viking_read(uri: String, _level: Option) -> Result { let storage = get_storage().await?; let entry = storage .get(&uri) .await .map_err(|e| format!("Failed to read memory: {}", e))?; match entry { Some(e) => Ok(e.content), None => Err(format!("Memory not found: {}", uri)), } } /// Remove a memory #[tauri::command] pub async fn viking_remove(uri: String) -> Result<(), String> { let storage = get_storage().await?; storage .delete(&uri) .await .map_err(|e| format!("Failed to remove memory: {}", e))?; Ok(()) } /// Get memory tree #[tauri::command] pub async fn viking_tree(path: String, _depth: Option) -> Result { let storage = get_storage().await?; let entries = storage .find_by_prefix(&path) .await .map_err(|e| format!("Failed to get tree: {}", e))?; // Build a simple tree structure let mut tree = serde_json::Map::new(); for entry in entries { let parts: Vec<&str> = entry.uri.split('/').collect(); let mut current = &mut tree; for part in &parts[..parts.len().saturating_sub(1)] { if !current.contains_key(*part) { current.insert( (*part).to_string(), serde_json::json!({}), ); } current = current .get_mut(*part) .and_then(|v| v.as_object_mut()) .unwrap(); } if let Some(last) = parts.last() { current.insert( (*last).to_string(), serde_json::json!({ "type": entry.memory_type.to_string(), "importance": entry.importance, "access_count": entry.access_count, }), ); } } Ok(serde_json::Value::Object(tree)) } /// Inject memories into prompt (for agent loop integration) #[tauri::command] pub async fn viking_inject_prompt( agent_id: String, base_prompt: String, user_input: String, max_tokens: Option, ) -> Result { let storage = get_storage().await?; // Retrieve relevant memories let options = FindOptions { scope: Some(format!("agent://{}", agent_id)), limit: Some(10), min_similarity: Some(0.3), }; let entries = storage .find(&user_input, options) .await .map_err(|e| format!("Failed to retrieve memories: {}", e))?; // Convert to RetrievalResult let mut result = RetrievalResult::default(); for entry in entries { match entry.memory_type { MemoryType::Preference => result.preferences.push(entry), MemoryType::Knowledge => result.knowledge.push(entry), MemoryType::Experience => result.experience.push(entry), MemoryType::Session => {} // Skip session memories } } // Calculate tokens result.total_tokens = result.calculate_tokens(); // Apply token budget let budget = max_tokens.unwrap_or(500); if result.total_tokens > budget { // Truncate by priority: preferences > knowledge > experience while result.total_tokens > budget && !result.experience.is_empty() { result.experience.pop(); result.total_tokens = result.calculate_tokens(); } while result.total_tokens > budget && !result.knowledge.is_empty() { result.knowledge.pop(); result.total_tokens = result.calculate_tokens(); } while result.total_tokens > budget && !result.preferences.is_empty() { result.preferences.pop(); result.total_tokens = result.calculate_tokens(); } } // Inject into prompt let injector = PromptInjector::new(); Ok(injector.inject_with_format(&base_prompt, &result)) } // === Helper Functions === /// Parse URI to extract components fn parse_uri(uri: &str) -> Result<(String, MemoryType, String), String> { // Expected format: agent://{agent_id}/{type}/{category} let without_prefix = uri .strip_prefix("agent://") .ok_or_else(|| format!("Invalid URI format: {}", uri))?; let parts: Vec<&str> = without_prefix.splitn(3, '/').collect(); if parts.len() < 3 { return Err(format!("Invalid URI format, expected agent://{{agent_id}}/{{type}}/{{category}}: {}", uri)); } let agent_id = parts[0].to_string(); let memory_type = MemoryType::parse(parts[1]); let category = parts[2].to_string(); Ok((agent_id, memory_type, category)) } // === Tests === #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_parse_uri() { let (agent_id, memory_type, category) = parse_uri("agent://test-agent/preferences/style").unwrap(); assert_eq!(agent_id, "test-agent"); assert_eq!(memory_type, MemoryType::Preference); assert_eq!(category, "style"); } #[test] fn test_invalid_uri() { assert!(parse_uri("invalid-uri").is_err()); assert!(parse_uri("agent://only-agent").is_err()); } }