memvid_cli/commands/

enrich.rs

//! Enrichment command handler for extracting memory cards from frames.
//!
//! The `enrich` command runs enrichment engines over MV2 frames to extract
//! structured memory cards (facts, preferences, events, etc.).

use std::path::PathBuf;

use anyhow::Result;
#[cfg(feature = "llama-cpp")]
use anyhow::bail;
use clap::{Args, ValueEnum};
use memvid_core::{EnrichmentEngine, Memvid, RulesEngine};
use serde::Serialize;

#[cfg(feature = "llama-cpp")]
use crate::commands::{default_enrichment_model, get_installed_model_path, LlmModel};
use crate::config::CliConfig;
#[cfg(feature = "candle-llm")]
use crate::enrich::CandlePhiEngine;
#[cfg(feature = "llama-cpp")]
use crate::enrich::LlmEngine;
use crate::enrich::OpenAiEngine;

/// Engine type for enrichment
#[derive(Debug, Clone, Copy, ValueEnum, Default)]
pub enum EnrichEngine {
    /// Rules-based extraction using regex patterns (fast, no models)
    #[default]
    Rules,
    /// LLM-based extraction with Phi-3.5 Mini via llama.cpp (requires model installation)
    #[cfg(feature = "llama-cpp")]
    Llm,
    /// Candle-based Phi-3 extraction (downloads from Hugging Face)
    #[cfg(feature = "candle-llm")]
    Candle,
    /// OpenAI API-based extraction with GPT-4o-mini (requires OPENAI_API_KEY)
    Openai,
}

/// Arguments for the `enrich` subcommand
#[derive(Args)]
pub struct EnrichArgs {
    /// Path to the `.mv2` file
    #[arg(value_name = "FILE", value_parser = clap::value_parser!(PathBuf))]
    pub file: PathBuf,

    /// Enrichment engine to use
    #[arg(long, value_enum, default_value_t = EnrichEngine::Rules)]
    pub engine: EnrichEngine,

    /// Only process frames that haven't been enriched yet (default)
    #[arg(long, default_value_t = true)]
    pub incremental: bool,

    /// Re-enrich all frames, ignoring previous enrichment records
    #[arg(long, conflicts_with = "incremental")]
    pub force: bool,

    /// Output results as JSON
    #[arg(long)]
    pub json: bool,

    /// Show extracted memory cards
    #[arg(long)]
    pub verbose: bool,
}

/// Result of enrichment for JSON output
#[derive(Debug, Serialize)]
pub struct EnrichResult {
    pub engine: String,
    pub version: String,
    pub frames_processed: usize,
    pub cards_extracted: usize,
    pub total_cards: usize,
    pub total_entities: usize,
}

/// Handle the `enrich` command
#[allow(unused_variables)]
pub fn handle_enrich(config: &CliConfig, args: EnrichArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Get initial stats
    let initial_stats = mem.memories_stats();

    // If force mode, clear existing memories first
    if args.force {
        mem.clear_memories();
    }

    // Run the selected engine
    let (engine_kind, engine_version, frames, cards) = match args.engine {
        EnrichEngine::Rules => {
            let engine = RulesEngine::new();
            let kind = engine.kind().to_string();
            let version = engine.version().to_string();
            let (frames, cards) = mem.run_enrichment(&engine)?;
            (kind, version, frames, cards)
        }
        #[cfg(feature = "llama-cpp")]
        EnrichEngine::Llm => {
            // Check if model is installed
            let model = default_enrichment_model();
            let model_path = match get_installed_model_path(config, model) {
                Some(path) => path,
                None => {
                    bail!(
                        "LLM model not installed. Run `memvid models install {}` first.",
                        match model {
                            LlmModel::Phi35Mini => "phi-3.5-mini",
                            LlmModel::Phi35MiniQ8 => "phi-3.5-mini-q8",
                        }
                    );
                }
            };

            // Create and initialize the LLM engine
            let mut engine = LlmEngine::new(model_path);
            eprintln!("Loading LLM model...");
            engine.init()?;

            let kind = engine.kind().to_string();
            let version = engine.version().to_string();
            let (frames, cards) = mem.run_enrichment(&engine)?;
            (kind, version, frames, cards)
        }
        #[cfg(feature = "candle-llm")]
        EnrichEngine::Candle => {
            // Create and initialize the Candle Phi-3 engine
            // Uses Q4 quantized GGUF (~2.4GB) stored in ~/.memvid/models/llm/phi-3-mini-q4/
            eprintln!("Loading Phi-3-mini Q4 model via Candle (first run downloads ~2.4GB to ~/.memvid/models/llm/)...");
            let mut engine = CandlePhiEngine::from_memvid_models(config.models_dir.clone());
            engine.init()?;

            let kind = engine.kind().to_string();
            let version = engine.version().to_string();
            let (frames, cards) = mem.run_enrichment(&engine)?;
            (kind, version, frames, cards)
        }
        EnrichEngine::Openai => {
            // Create and initialize the OpenAI engine with parallel batch support
            eprintln!("Using OpenAI GPT-4o-mini for enrichment (parallel mode)...");
            let mut engine = OpenAiEngine::new();
            engine.init()?;

            let kind = engine.kind().to_string();
            let version = engine.version().to_string();

            // Use parallel batch processing for OpenAI
            let (frames, cards) = run_openai_parallel(&mut mem, &engine)?;
            (kind, version, frames, cards)
        }
    };

    // Commit changes
    mem.commit()?;

    // Get final stats
    let final_stats = mem.memories_stats();

    if args.json {
        let result = EnrichResult {
            engine: engine_kind,
            version: engine_version,
            frames_processed: frames,
            cards_extracted: cards,
            total_cards: final_stats.card_count,
            total_entities: final_stats.entity_count,
        };
        println!("{}", serde_json::to_string_pretty(&result)?);
    } else {
        println!("Enrichment complete:");
        println!("  Engine: {} v{}", engine_kind, engine_version);
        println!("  Frames processed: {}", frames);
        println!("  Cards extracted: {}", cards);
        println!(
            "  Total cards: {} (+{})",
            final_stats.card_count,
            final_stats
                .card_count
                .saturating_sub(initial_stats.card_count)
        );
        println!("  Entities: {}", final_stats.entity_count);

        if args.verbose && cards > 0 {
            println!("\nExtracted memory cards:");
            for entity in mem.memory_entities() {
                println!("  {}:", entity);
                for card in mem.get_entity_memories(&entity) {
                    println!("    - {}: {} = \"{}\"", card.kind, card.slot, card.value);
                }
            }
        }
    }

    Ok(())
}

/// Run OpenAI enrichment with parallel batch processing.
///
/// This gathers all unenriched frames, sends them to OpenAI in parallel (20 concurrent requests),
/// and stores the resulting memory cards. This is ~20x faster than sequential processing.
fn run_openai_parallel(
    mem: &mut memvid_core::Memvid,
    engine: &OpenAiEngine,
) -> Result<(usize, usize)> {
    use memvid_core::enrich::EnrichmentContext;
    use memvid_core::EnrichmentEngine;

    let kind = engine.kind();
    let version = engine.version();

    // Get all unenriched frames
    let unenriched = mem.get_unenriched_frames(kind, version);
    let total_frames = unenriched.len();

    if total_frames == 0 {
        eprintln!("No unenriched frames found.");
        return Ok((0, 0));
    }

    eprintln!(
        "Gathering {} frames for parallel enrichment...",
        total_frames
    );

    // Build enrichment contexts for all frames
    let mut contexts = Vec::with_capacity(total_frames);
    for frame_id in &unenriched {
        let frame = match mem.frame_by_id(*frame_id) {
            Ok(f) => f,
            Err(_) => continue,
        };

        // Get frame content via preview (text extraction)
        let text = match mem.frame_preview_by_id(*frame_id) {
            Ok(t) => t,
            Err(_) => continue,
        };

        let uri = frame
            .uri
            .clone()
            .unwrap_or_else(|| format!("mv2://frame/{}", frame_id));
        let metadata_json = frame
            .metadata
            .as_ref()
            .and_then(|m| serde_json::to_string(m).ok());

        let ctx = EnrichmentContext::new(
            *frame_id,
            uri,
            text,
            frame.title.clone(),
            frame.timestamp,
            metadata_json,
        );

        contexts.push(ctx);
    }

    eprintln!(
        "Starting parallel enrichment of {} frames with 20 workers...",
        contexts.len()
    );

    // Run parallel batch enrichment
    let results = engine.enrich_batch(contexts);

    // Store results back to MV2
    let mut total_cards = 0;
    for (frame_id, cards) in results {
        let card_count = cards.len();

        // Store cards
        let card_ids = if !cards.is_empty() {
            mem.put_memory_cards(cards)?
        } else {
            Vec::new()
        };

        // Record enrichment
        mem.record_enrichment(frame_id, kind, version, card_ids)?;

        total_cards += card_count;
    }

    Ok((total_frames, total_cards))
}

/// Handle the `memories` subcommand (view memory cards)
#[derive(Args)]
pub struct MemoriesArgs {
    /// Path to the `.mv2` file
    #[arg(value_name = "FILE", value_parser = clap::value_parser!(PathBuf))]
    pub file: PathBuf,

    /// Filter by entity
    #[arg(long)]
    pub entity: Option<String>,

    /// Filter by slot
    #[arg(long)]
    pub slot: Option<String>,

    /// Output as JSON
    #[arg(long)]
    pub json: bool,
}

/// Memory card output for JSON serialization
#[derive(Debug, Serialize)]
pub struct MemoryOutput {
    pub id: u64,
    pub kind: String,
    pub entity: String,
    pub slot: String,
    pub value: String,
    pub polarity: Option<String>,
    pub document_date: Option<i64>,
    pub source_frame_id: u64,
}

pub fn handle_memories(_config: &CliConfig, args: MemoriesArgs) -> Result<()> {
    let mem = Memvid::open(&args.file)?;

    let stats = mem.memories_stats();

    if args.json {
        let mut cards: Vec<MemoryOutput> = Vec::new();

        if let Some(entity) = &args.entity {
            if let Some(slot) = &args.slot {
                // Specific entity:slot
                if let Some(card) = mem.get_current_memory(entity, slot) {
                    cards.push(card_to_output(card));
                }
            } else {
                // All cards for entity
                for card in mem.get_entity_memories(entity) {
                    cards.push(card_to_output(card));
                }
            }
        } else {
            // All entities
            for entity in mem.memory_entities() {
                for card in mem.get_entity_memories(&entity) {
                    cards.push(card_to_output(card));
                }
            }
        }

        println!("{}", serde_json::to_string_pretty(&cards)?);
    } else {
        println!(
            "Memory cards: {} total, {} entities",
            stats.card_count, stats.entity_count
        );
        println!();

        if let Some(entity) = &args.entity {
            if let Some(slot) = &args.slot {
                // Specific entity:slot
                if let Some(card) = mem.get_current_memory(entity, slot) {
                    println!("{}:{} = \"{}\"", entity, slot, card.value);
                } else {
                    println!("No memory found for {}:{}", entity, slot);
                }
            } else {
                // All cards for entity
                println!("{}:", entity);
                for card in mem.get_entity_memories(entity) {
                    println!("  {}: {} = \"{}\"", card.kind, card.slot, card.value);
                }
            }
        } else {
            // All entities
            for entity in mem.memory_entities() {
                println!("{}:", entity);
                for card in mem.get_entity_memories(&entity) {
                    let polarity = card
                        .polarity
                        .as_ref()
                        .map(|p| format!(" [{}]", p))
                        .unwrap_or_default();
                    println!(
                        "  {}: {} = \"{}\"{}",
                        card.kind, card.slot, card.value, polarity
                    );
                }
                println!();
            }
        }
    }

    Ok(())
}

fn card_to_output(card: &memvid_core::MemoryCard) -> MemoryOutput {
    MemoryOutput {
        id: card.id,
        kind: card.kind.to_string(),
        entity: card.entity.clone(),
        slot: card.slot.clone(),
        value: card.value.clone(),
        polarity: card.polarity.as_ref().map(|p| p.to_string()),
        document_date: card.document_date,
        source_frame_id: card.source_frame_id,
    }
}