zclaw_openfang/desktop/src-tauri/src/memory/extractor.rs

//! 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<String>,
    pub importance: u8, // 1-10 scale
    pub suggested_uri: String,
    pub reasoning: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExtractionResult {
    pub memories: Vec<ExtractedMemory>,
    pub summary: String,
    pub tokens_saved: Option<u32>,
    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<String>,
}

// === Session Extractor ===

pub struct SessionExtractor {
    config: ExtractionConfig,
    llm_endpoint: Option<String>,
    api_key: Option<String>,
}

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<ExtractionResult, String> {
        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<ExtractedMemory> = 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::<Vec<_>>()
            .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<String, String> {
        // 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<String, String> {
        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<String, String> {
        // Simple pattern matching for common memory patterns
        let mut memories: Vec<ExtractedMemory> = 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<Vec<ExtractedMemory>, 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::Value> = serde_json::from_str(json_str)
            .unwrap_or_default();

        let memories: Vec<ExtractedMemory> = 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<ExtractedMemory> {
        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<ChatMessage>,
    agent_id: String,
) -> Result<ExtractionResult, String> {
    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<ChatMessage>,
    agent_id: String,
    llm_endpoint: Option<String>,
    llm_api_key: Option<String>,
) -> Result<ExtractionResult, String> {
    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/"));
    }
}