//! Session Memory Extractor //! //! Extracts structured memories from conversation sessions using LLM analysis. //! This supplements OpenViking CLI which lacks built-in memory extraction. //! //! Categories: //! - user_preference: User's stated preferences and settings //! - user_fact: Facts about the user (name, role, projects, etc.) //! - agent_lesson: Lessons learned by the agent from interactions //! - agent_pattern: Recurring patterns the agent should remember //! - task: Task-related information for follow-up //! //! Note: Some fields and methods are reserved for future LLM-powered extraction // NOTE: #[tauri::command] functions are registered via invoke_handler! at runtime. // Module-level allow required for Tauri-commanded functions and internal types. #![allow(dead_code)] use serde::{Deserialize, Serialize}; // === Types === #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum MemoryCategory { UserPreference, UserFact, AgentLesson, AgentPattern, Task, } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExtractedMemory { pub category: MemoryCategory, pub content: String, pub tags: Vec, pub importance: u8, // 1-10 scale pub suggested_uri: String, pub reasoning: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExtractionResult { pub memories: Vec, pub summary: String, pub tokens_saved: Option, pub extraction_time_ms: u64, } #[derive(Debug, Clone)] pub struct ExtractionConfig { /// Maximum memories to extract per session pub max_memories: usize, /// Minimum importance threshold (1-10) pub min_importance: u8, /// Whether to include reasoning in output pub include_reasoning: bool, /// Agent ID for URI generation pub agent_id: String, } impl Default for ExtractionConfig { fn default() -> Self { Self { max_memories: 10, min_importance: 5, include_reasoning: true, agent_id: "zclaw-main".to_string(), } } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { pub role: String, pub content: String, pub timestamp: Option, } // === Session Extractor === pub struct SessionExtractor { config: ExtractionConfig, llm_endpoint: Option, api_key: Option, } impl SessionExtractor { pub fn new(config: ExtractionConfig) -> Self { Self { config, llm_endpoint: None, api_key: None, } } /// Configure LLM endpoint for extraction pub fn with_llm(mut self, endpoint: String, api_key: String) -> Self { self.llm_endpoint = Some(endpoint); self.api_key = Some(api_key); self } /// Extract memories from a conversation session pub async fn extract(&self, messages: &[ChatMessage]) -> Result { let start_time = std::time::Instant::now(); // Build extraction prompt let prompt = self.build_extraction_prompt(messages); // Call LLM for extraction let response = self.call_llm(&prompt).await?; // Parse LLM response into structured memories let memories = self.parse_extraction(&response)?; // Filter by importance and limit let filtered: Vec = memories .into_iter() .filter(|m| m.importance >= self.config.min_importance) .take(self.config.max_memories) .collect(); // Generate session summary let summary = self.generate_summary(&filtered); let elapsed = start_time.elapsed().as_millis() as u64; Ok(ExtractionResult { tokens_saved: Some(self.estimate_tokens_saved(messages, &summary)), memories: filtered, summary, extraction_time_ms: elapsed, }) } /// Build the extraction prompt for the LLM fn build_extraction_prompt(&self, messages: &[ChatMessage]) -> String { let conversation = messages .iter() .map(|m| format!("[{}]: {}", m.role, m.content)) .collect::>() .join("\n\n"); format!( r#"Analyze the following conversation and extract structured memories. Focus on information that would be useful for future interactions. ## Conversation {} ## Extraction Instructions Extract memories in these categories: - user_preference: User's stated preferences (UI preferences, workflow preferences, tool choices) - user_fact: Facts about the user (name, role, projects, skills, constraints) - agent_lesson: Lessons the agent learned (what worked, what didn't, corrections needed) - agent_pattern: Recurring patterns to remember (common workflows, frequent requests) - task: Tasks or follow-ups mentioned (todos, pending work, deadlines) For each memory, provide: 1. category: One of the above categories 2. content: The actual memory content (concise, actionable) 3. tags: 2-5 relevant tags for retrieval 4. importance: 1-10 scale (10 = critical, 1 = trivial) 5. reasoning: Brief explanation of why this is worth remembering Output as JSON array: ```json [ {{ "category": "user_preference", "content": "...", "tags": ["tag1", "tag2"], "importance": 7, "reasoning": "..." }} ] ``` If no significant memories found, return empty array: []"#, conversation ) } /// Call LLM for extraction async fn call_llm(&self, prompt: &str) -> Result { // If LLM endpoint is configured, use it if let (Some(endpoint), Some(api_key)) = (&self.llm_endpoint, &self.api_key) { return self.call_llm_api(endpoint, api_key, prompt).await; } // Otherwise, use rule-based extraction as fallback self.rule_based_extraction(prompt) } /// Call external LLM API (doubao, OpenAI, etc.) async fn call_llm_api( &self, endpoint: &str, api_key: &str, prompt: &str, ) -> Result { let client = reqwest::Client::new(); let response = client .post(endpoint) .header("Authorization", format!("Bearer {}", api_key)) .header("Content-Type", "application/json") .json(&serde_json::json!({ "model": "doubao-pro-32k", "messages": [ {"role": "user", "content": prompt} ], "temperature": 0.3, "max_tokens": 2000 })) .send() .await .map_err(|e| format!("LLM API request failed: {}", e))?; if !response.status().is_success() { return Err(format!("LLM API error: {}", response.status())); } let json: serde_json::Value = response .json() .await .map_err(|e| format!("Failed to parse LLM response: {}", e))?; // Extract content from response (adjust based on API format) let content = json .get("choices") .and_then(|c| c.get(0)) .and_then(|c| c.get("message")) .and_then(|m| m.get("content")) .and_then(|c| c.as_str()) .ok_or("Invalid LLM response format")? .to_string(); Ok(content) } /// Rule-based extraction as fallback when LLM is not available fn rule_based_extraction(&self, prompt: &str) -> Result { // Simple pattern matching for common memory patterns let mut memories: Vec = Vec::new(); // Pattern: User preferences let pref_patterns = [ (r"I prefer (.+)", "user_preference"), (r"My preference is (.+)", "user_preference"), (r"I like (.+)", "user_preference"), (r"I don't like (.+)", "user_preference"), ]; // Pattern: User facts let fact_patterns = [ (r"My name is (.+)", "user_fact"), (r"I work on (.+)", "user_fact"), (r"I'm a (.+)", "user_fact"), (r"My project is (.+)", "user_fact"), ]; // Extract using regex (simplified implementation) for (pattern, category) in pref_patterns.iter().chain(fact_patterns.iter()) { if let Ok(re) = regex::Regex::new(pattern) { for cap in re.captures_iter(prompt) { if let Some(content) = cap.get(1) { let memory = ExtractedMemory { category: if *category == "user_preference" { MemoryCategory::UserPreference } else { MemoryCategory::UserFact }, content: content.as_str().to_string(), tags: vec!["auto-extracted".to_string()], importance: 6, suggested_uri: format!( "viking://user/memories/{}/{}", category, chrono::Utc::now().timestamp_millis() ), reasoning: Some("Extracted via rule-based pattern matching".to_string()), }; memories.push(memory); } } } } // Return as JSON serde_json::to_string_pretty(&memories) .map_err(|e| format!("Failed to serialize memories: {}", e)) } /// Parse LLM response into structured memories fn parse_extraction(&self, response: &str) -> Result, String> { // Try to extract JSON from the response let json_start = response.find('[').unwrap_or(0); let json_end = response.rfind(']').map(|i| i + 1).unwrap_or(response.len()); let json_str = &response[json_start..json_end]; // Parse JSON let raw_memories: Vec = serde_json::from_str(json_str) .unwrap_or_default(); let memories: Vec = raw_memories .into_iter() .filter_map(|m| self.parse_memory(&m)) .collect(); Ok(memories) } /// Parse a single memory from JSON fn parse_memory(&self, value: &serde_json::Value) -> Option { let category_str = value.get("category")?.as_str()?; let category = match category_str { "user_preference" => MemoryCategory::UserPreference, "user_fact" => MemoryCategory::UserFact, "agent_lesson" => MemoryCategory::AgentLesson, "agent_pattern" => MemoryCategory::AgentPattern, "task" => MemoryCategory::Task, _ => return None, }; let content = value.get("content")?.as_str()?.to_string(); let tags = value .get("tags") .and_then(|t| t.as_array()) .map(|arr| { arr.iter() .filter_map(|v| v.as_str().map(String::from)) .collect() }) .unwrap_or_default(); let importance = value .get("importance") .and_then(|v| v.as_u64()) .unwrap_or(5) as u8; let reasoning = value .get("reasoning") .and_then(|v| v.as_str()) .map(String::from); // Generate URI based on category let suggested_uri = self.generate_uri(&category, &content); Some(ExtractedMemory { category, content, tags, importance, suggested_uri, reasoning, }) } /// Generate a URI for the memory fn generate_uri(&self, category: &MemoryCategory, content: &str) -> String { let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis()) .unwrap_or(0); let content_hash = &content[..content.len().min(20)] .to_lowercase() .replace(' ', "_") .replace(|c: char| !c.is_alphanumeric() && c != '_', ""); match category { MemoryCategory::UserPreference => { format!("viking://user/memories/preferences/{}_{}", content_hash, timestamp) } MemoryCategory::UserFact => { format!("viking://user/memories/facts/{}_{}", content_hash, timestamp) } MemoryCategory::AgentLesson => { format!( "viking://agent/{}/memories/lessons/{}_{}", self.config.agent_id, content_hash, timestamp ) } MemoryCategory::AgentPattern => { format!( "viking://agent/{}/memories/patterns/{}_{}", self.config.agent_id, content_hash, timestamp ) } MemoryCategory::Task => { format!( "viking://agent/{}/tasks/{}_{}", self.config.agent_id, content_hash, timestamp ) } } } /// Generate a summary of extracted memories fn generate_summary(&self, memories: &[ExtractedMemory]) -> String { if memories.is_empty() { return "No significant memories extracted from this session.".to_string(); } let mut summary_parts = Vec::new(); let user_prefs = memories .iter() .filter(|m| matches!(m.category, MemoryCategory::UserPreference)) .count(); if user_prefs > 0 { summary_parts.push(format!("{} user preferences", user_prefs)); } let user_facts = memories .iter() .filter(|m| matches!(m.category, MemoryCategory::UserFact)) .count(); if user_facts > 0 { summary_parts.push(format!("{} user facts", user_facts)); } let lessons = memories .iter() .filter(|m| matches!(m.category, MemoryCategory::AgentLesson)) .count(); if lessons > 0 { summary_parts.push(format!("{} agent lessons", lessons)); } let patterns = memories .iter() .filter(|m| matches!(m.category, MemoryCategory::AgentPattern)) .count(); if patterns > 0 { summary_parts.push(format!("{} patterns", patterns)); } let tasks = memories .iter() .filter(|m| matches!(m.category, MemoryCategory::Task)) .count(); if tasks > 0 { summary_parts.push(format!("{} tasks", tasks)); } format!( "Extracted {} memories: {}.", memories.len(), summary_parts.join(", ") ) } /// Estimate tokens saved by extraction fn estimate_tokens_saved(&self, messages: &[ChatMessage], summary: &str) -> u32 { // Rough estimation: original messages vs summary let original_tokens: u32 = messages .iter() .map(|m| (m.content.len() as f32 * 0.4) as u32) .sum(); let summary_tokens = (summary.len() as f32 * 0.4) as u32; original_tokens.saturating_sub(summary_tokens) } } // === Tauri Commands === // @reserved: memory extraction // @connected #[tauri::command] pub async fn extract_session_memories( messages: Vec, agent_id: String, ) -> Result { let config = ExtractionConfig { agent_id, ..Default::default() }; let extractor = SessionExtractor::new(config); extractor.extract(&messages).await } /// Extract memories from session and store to SqliteStorage /// This combines extraction and storage in one command // @reserved: memory extraction and storage // @connected #[tauri::command] pub async fn extract_and_store_memories( messages: Vec, agent_id: String, llm_endpoint: Option, llm_api_key: Option, ) -> Result { use zclaw_growth::{MemoryEntry, MemoryType, VikingStorage}; let start_time = std::time::Instant::now(); // 1. Extract memories let config = ExtractionConfig { agent_id: agent_id.clone(), ..Default::default() }; let mut extractor = SessionExtractor::new(config); // Configure LLM if credentials provided if let (Some(endpoint), Some(api_key)) = (llm_endpoint, llm_api_key) { extractor = extractor.with_llm(endpoint, api_key); } let extraction_result = extractor.extract(&messages).await?; // 2. Get storage instance let storage = crate::viking_commands::get_storage() .await .map_err(|e| format!("Storage not available: {}", e))?; // 3. Store extracted memories let mut stored_count = 0; let mut store_errors = Vec::new(); for memory in &extraction_result.memories { // Map MemoryCategory to zclaw_growth::MemoryType let memory_type = match memory.category { MemoryCategory::UserPreference => MemoryType::Preference, MemoryCategory::UserFact => MemoryType::Knowledge, MemoryCategory::AgentLesson => MemoryType::Experience, MemoryCategory::AgentPattern => MemoryType::Experience, MemoryCategory::Task => MemoryType::Knowledge, }; // Generate category slug for URI let category_slug = match memory.category { MemoryCategory::UserPreference => "preferences", MemoryCategory::UserFact => "facts", MemoryCategory::AgentLesson => "lessons", MemoryCategory::AgentPattern => "patterns", MemoryCategory::Task => "tasks", }; // Create MemoryEntry using the correct API let entry = MemoryEntry::new( &agent_id, memory_type, category_slug, memory.content.clone(), ) .with_keywords(memory.tags.clone()) .with_importance(memory.importance); // Store to SqliteStorage let entry_uri = entry.uri.clone(); match storage.store(&entry).await { Ok(_) => stored_count += 1, Err(e) => { store_errors.push(format!("Failed to store {}: {}", memory.category, e)); } } // Background: generate L0/L1 summaries if driver is configured if crate::summarizer_adapter::is_summary_driver_configured() { let storage_clone = storage.clone(); let summary_entry = entry.clone(); tokio::spawn(async move { if let Some(driver) = crate::summarizer_adapter::get_summary_driver() { let (overview, abstract_summary) = zclaw_growth::summarizer::generate_summaries(driver.as_ref(), &summary_entry).await; if overview.is_some() || abstract_summary.is_some() { let updated = MemoryEntry { overview, abstract_summary, ..summary_entry }; if let Err(e) = storage_clone.store(&updated).await { tracing::debug!( "[extract_and_store] Failed to update summaries for {}: {}", entry_uri, e ); } } } }); } } let elapsed = start_time.elapsed().as_millis() as u64; // Log any storage errors if !store_errors.is_empty() { tracing::warn!( "[extract_and_store] {} memories stored, {} errors: {}", stored_count, store_errors.len(), store_errors.join("; ") ); } tracing::info!( "[extract_and_store] Extracted {} memories, stored {} in {}ms", extraction_result.memories.len(), stored_count, elapsed ); // Return updated result with storage info Ok(ExtractionResult { memories: extraction_result.memories, summary: format!( "{} (Stored: {})", extraction_result.summary, stored_count ), tokens_saved: extraction_result.tokens_saved, extraction_time_ms: elapsed, }) } impl std::fmt::Display for MemoryCategory { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { MemoryCategory::UserPreference => write!(f, "user_preference"), MemoryCategory::UserFact => write!(f, "user_fact"), MemoryCategory::AgentLesson => write!(f, "agent_lesson"), MemoryCategory::AgentPattern => write!(f, "agent_pattern"), MemoryCategory::Task => write!(f, "task"), } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_extraction_config_default() { let config = ExtractionConfig::default(); assert_eq!(config.max_memories, 10); assert_eq!(config.min_importance, 5); } #[test] fn test_uri_generation() { let config = ExtractionConfig::default(); let extractor = SessionExtractor::new(config); let uri = extractor.generate_uri( &MemoryCategory::UserPreference, "dark mode enabled" ); assert!(uri.starts_with("viking://user/memories/preferences/")); } }