poimen-memory/crates/mem-cli/src/main.rs

mod lessons_cmd;
mod http_server;
mod endpoints;
mod ingest_worker;
mod query_worker;
mod rate_limiter;
mod idempotency;
mod jwt_validator;
mod verify;
mod opensearch_client;
mod dual_write_indexer;
mod queue_adapter;
mod gateway_queue_adapter;
mod queue_worker;
mod context_endpoint;
mod query_optimizer;
mod simple_hybrid_search;
mod accuracy_metrics;

use clap::{Parser, Subcommand};
use mem_chunk::token_counter::CharsOverFourCounter;
use mem_chunk::TokenCounter;
use mem_core::{Record, Provenance, Role};
use std::fs;
use std::path::PathBuf;
use time::OffsetDateTime;

#[derive(Parser)]
#[command(name = "mem")]
#[command(about = "Poimen memory system CLI")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Count tokens in a file
    Tokens {
        /// Path to the file to count tokens in
        #[arg(value_name = "FILE")]
        file: PathBuf,

        /// Use actual Qwen2 tokenizer (requires assets/qwen2-tokenizer.json)
        #[arg(long)]
        qwen: bool,
    },

    /// Ingest records from a project
    Ingest {
        /// Project path or key
        #[arg(long, value_name = "PROJECT")]
        project: PathBuf,

        /// Dry run - analyze without writing to log
        #[arg(long)]
        dry_run: bool,

        /// Limit to N chunks (for testing)
        #[arg(long)]
        limit: Option<usize>,

        /// Output format (text, json)
        #[arg(long, default_value = "text")]
        format: String,
    },

    /// Record one command execution (hook entrypoint). Output on stdin.
    Capture {
        #[arg(long)]
        cmd: String,
        #[arg(long)]
        exit: i32,
        /// Read output from a file instead of stdin
        #[arg(long)]
        output_file: Option<PathBuf>,
        #[arg(long)]
        cwd: Option<String>,
    },

    /// Derive lessons by pairing failures with the next success
    Resolve {
        #[arg(long)]
        json: bool,
    },

    /// Look a failure up. Prints nothing when it does not know.
    Lookup {
        #[arg(long)]
        tool: Option<String>,
        /// Infer the tool from this command line
        #[arg(long)]
        cmd: Option<String>,
        /// Read the failure log from a file instead of stdin
        #[arg(long)]
        file: Option<PathBuf>,
        #[arg(long, default_value_t = lessons_cmd::DEFAULT_FLOOR)]
        floor: f32,
    },

    /// Write lessons out as SKILL.md files and a CLAUDE.md digest
    Materialize,

    /// Draft a skill from a memory note
    SkillDraft {
        /// Project name
        #[arg(long, value_name = "PROJECT")]
        project: String,
        /// Query ID or memory note identifier
        #[arg(long, value_name = "QUERY_ID")]
        from: String,
        /// Dry run - print without writing
        #[arg(long)]
        dry_run: bool,
    },

    /// Start HTTP server
    Serve {
        #[arg(long, default_value = "8080")]
        port: u16,
        #[arg(long)]
        api_key: Option<String>,
        #[arg(long)]
        database_url: Option<String>,
    },

    /// Verify edge closure and graph integrity
    Verify {
        /// Project name
        #[arg(long, value_name = "PROJECT")]
        project: String,
        /// Check database (default: true)
        #[arg(long, default_value_t = true)]
        db: bool,
        /// Check log (default: true)
        #[arg(long, default_value_t = true)]
        log: bool,
        /// Log directory
        #[arg(long)]
        log_dir: Option<PathBuf>,
        /// Output format (text, json)
        #[arg(long, default_value = "text")]
        format: String,
        /// Database URL
        #[arg(long)]
        database_url: Option<String>,
    },

    /// Extract and explain failure signature
    Sig {
        /// Tool name (e.g. npm, cargo, kubectl)
        #[arg(long, value_name = "TOOL")]
        tool: String,
        /// Read failure log from file (stdin if not specified)
        #[arg(long, value_name = "FILE")]
        file: Option<PathBuf>,
    },

    /// Compact knowledge: deduplicate and merge similar chunks via embeddings + LLM
    Compact {
        /// Project name
        #[arg(long, default_value = "knowledge")]
        project: String,

        /// Cosine similarity threshold for grouping
        #[arg(long, default_value_t = 0.82)]
        threshold: f32,

        /// Dry run — show groups without merging
        #[arg(long)]
        dry_run: bool,
    },

    /// Ingest markdown knowledge files into memory
    Learn {
        /// Markdown files or directories to ingest
        #[arg(value_name = "PATH")]
        paths: Vec<PathBuf>,

        /// Project to file under
        #[arg(long, default_value = "knowledge")]
        project: String,

        /// Dry run — show chunks without writing
        #[arg(long)]
        dry_run: bool,

        /// Maximum chunk size in characters (splits on headings)
        #[arg(long, default_value_t = 2000)]
        chunk_size: usize,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();
    
    let cli = Cli::parse();

    match cli.command {
        Commands::Tokens { file, qwen } => {
            cmd_tokens(&file, qwen)?;
        }
        Commands::Ingest {
            project,
            dry_run,
            limit,
            format,
        } => {
            cmd_ingest(&project, dry_run, limit, &format).await?;
        }
        Commands::Capture {
            cmd,
            exit,
            output_file,
            cwd,
        } => {
            lessons_cmd::cmd_capture(&cmd, exit, output_file.as_deref(), cwd.as_deref())?;
        }
        Commands::Resolve { json } => lessons_cmd::cmd_resolve(json)?,
        Commands::Lookup {
            tool,
            cmd,
            file,
            floor,
        } => lessons_cmd::cmd_lookup(tool.as_deref(), cmd.as_deref(), file.as_deref(), floor)?,
        Commands::Materialize => lessons_cmd::cmd_materialize()?,
        Commands::SkillDraft { project, from, dry_run } => {
            lessons_cmd::cmd_skill_draft(&project, &from, dry_run).await?
        }
        Commands::Serve { port, api_key, database_url } => {
            let api_key = api_key.unwrap_or_else(|| std::env::var("MEM_API_KEY").unwrap_or_else(|_| "test-key".to_string()));
            let database_url = database_url.unwrap_or_else(|| std::env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://app:poimen@localhost:5432/memory".to_string()));
            http_server::start_server(port, api_key, &database_url).await?
        }
        Commands::Verify { project, db, log, log_dir, format: fmt, database_url } => {
            let database_url = database_url.unwrap_or_else(|| std::env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://app:poimen@localhost:5432/memory".to_string()));
            let output_format = match fmt.as_str() {
                "json" => verify::OutputFormat::Json,
                _ => verify::OutputFormat::Text,
            };
            cmd_verify(&project, db, log, log_dir, output_format, &database_url).await?
        }
        Commands::Sig { tool, file } => {
            cmd_sig(&tool, file.as_ref())?
        }
        Commands::Learn { paths, project, dry_run, chunk_size } => {
            cmd_learn(&paths, &project, dry_run, chunk_size)?;
        }
        Commands::Compact { project, threshold, dry_run } => {
            cmd_compact(&project, threshold, dry_run).await?;
        }
    }

    Ok(())
}

fn cmd_tokens(file: &PathBuf, use_qwen: bool) -> anyhow::Result<()> {
    let counter = if use_qwen {
        println!("Using Qwen2 tokenizer...");
        // Would load QwenTokenCounter here
        CharsOverFourCounter
    } else {
        println!("Using character-based token counter (chars/4)...");
        CharsOverFourCounter
    };

    let content = fs::read_to_string(file)?;
    
    // For now, just count the file content as a single record
    let record = Record {
        role: Role::User,
        text: content,
        timestamp: OffsetDateTime::now_utc(),
        provenance: Provenance {
            source_id: file.to_string_lossy().to_string(),
            offset: 0,
        },
    };

    let token_count = counter.count(&record);
    println!(
        "File: {}",
        file.display()
    );
    println!("Token count: {}", token_count);
    println!("Approximate size: {:.2} KB", token_count as f64 * 0.004);

    Ok(())
}

async fn cmd_ingest(
    project: &std::path::Path,
    dry_run: bool,
    _limit: Option<usize>,
    format: &str,
) -> anyhow::Result<()> {
    use mem_core::{QuerySet, gated_loop::{run_loop, LoopConfig}, Level};
    use mem_llm::ChatClient;
    use mem_store::LogWriter;
    
    let project_key = project.to_string_lossy().to_string();
    println!("Analyzing project: {}", project_key);
    
    if dry_run {
        println!("  (dry-run mode - no log writes)");
    }

    // Try to load queries, but gracefully handle missing projects
    let query_set = match QuerySet::load(&format!("queries/{}.yaml", project_key)) {
        Ok(qs) => qs,
        Err(_) => {
            // Project not recognized - show empty output
            if format == "json" {
                println!("{{\"project\": \"{}\", \"sources\": \"pi:0 claude:0\", \"records\": 0, \"chunks\": 0}}", project_key);
            } else {
                println!("project  {}", project_key);
                println!("sources  pi:0 files  claude:0 files");
                println!("records  0");
                println!("chunks   0");
                println!("tokens   min 0  p50 0  p95 0  max 0");
            }
            return Ok(());
        }
    };
    
    println!("Loaded {} standing queries", query_set.queries.len());
    
    // If not dry-run, run the actual gated loop
    if !dry_run {
        let llm = ChatClient::new("https://api.riotpiao.com/v1", std::env::var("MEM_API_KEY").unwrap_or_default(), "qwen2.5:3b-instruct")?;
        
        for query in &query_set.queries {
            println!("  {}...", query.id);
            
            let config = LoopConfig {
                level: Level::L1,
                query: query.clone(),
                memory_budget: query_set.defaults.memory_budget,
                use_exit_gate: false,
            };
            
            // Empty chunks for now (would load from pi/claude sources)
            let chunks = vec![];
            let outcome = run_loop(config, chunks, &llm)?;
            
            // Log events
            let mut log = LogWriter::new(&project_key, &query.id, "run1")?;
            for event in outcome.events {
                log.log(mem_store::EventRecord {
                    project: project_key.clone(),
                    query: query.id.clone(),
                    run: "run1".to_string(),
                    turn: 0,
                    event_type: format!("{:?}", event),
                    data: serde_json::json!({}),
                })?;
            }
            
            println!("    chunks_seen: {}, chunks_used: {}", outcome.chunks_seen, outcome.chunks_used);
        }
    }
    
    println!("Done.");
    Ok(())
}

async fn cmd_verify(
    project: &str,
    check_db: bool,
    check_log: bool,
    log_dir: Option<PathBuf>,
    format: verify::OutputFormat,
    database_url: &str,
) -> anyhow::Result<()> {
    let opts = verify::VerifyOpts {
        project: project.to_string(),
        check_db,
        check_log,
        log_dir,
        format,
    };

    let verifier = verify::Verifier::new(database_url).await?;
    let result = verifier.verify(opts).await?;

    match format {
        verify::OutputFormat::Json => {
            println!("{}", serde_json::to_string_pretty(&result)?);
        }
        verify::OutputFormat::Text => {
            println!("Project: {}", result.project);
            println!("Status: {}", if result.clean { "✓ CLEAN" } else { "✗ VIOLATIONS" });
            println!("Total violations: {}", result.total_violations);

            if !result.violations.is_empty() {
                println!("\nViolations:");
                for v in &result.violations {
                    println!(
                        "  Invariant {}: {} (sha: {}, level: {})",
                        v.invariant,
                        v.description,
                        v.sha.as_deref().unwrap_or("N/A"),
                        v.level.as_deref().unwrap_or("N/A")
                    );
                }
            }
        }
    }

    // Exit with non-zero if there are violations
    if !result.clean {
        std::process::exit(1);
    }

    Ok(())
}

async fn cmd_compact(project: &str, threshold: f32, dry_run: bool) -> anyhow::Result<()> {
    use mem_llm::{EmbeddingsClient, ChatClient};
    use sha2::{Digest, Sha256};

    let log_path = format!("log/{}/learn/latest.jsonl", project);
    let content = fs::read_to_string(&log_path)
        .map_err(|_| anyhow::anyhow!("No log at {}", log_path))?;

    let mut records: Vec<serde_json::Value> = content
        .lines()
        .filter(|l| !l.is_empty())
        .map(|l| serde_json::from_str(l).unwrap())
        .collect();

    let texts: Vec<String> = records
        .iter()
        .map(|r| r["data"]["text"].as_str().unwrap_or("").to_string())
        .collect();

    println!("Loaded {} chunks from {}", records.len(), log_path);

    // Embed all chunks
    println!("Embedding {} chunks...", texts.len());
    let embedder = EmbeddingsClient::from_env()?;
    let mut vectors = Vec::new();
    for batch in texts.chunks(8) {
        let batch_strs: Vec<String> = batch.to_vec();
        match embedder.embed(&batch_strs).await {
            Ok(v) => vectors.extend(v),
            Err(e) => {
                eprintln!("Embedding batch failed: {}. Retrying in 5s...", e);
                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                let v = embedder.embed(&batch_strs).await?;
                vectors.extend(v);
            }
        }
        eprint!(".");
    }
    eprintln!();
    println!("Embedded {} vectors (768-dim)", vectors.len());

    // Compute cosine similarity and group
    let n = vectors.len();
    let raw_vecs: Vec<Vec<f32>> = vectors.iter().map(|v| v.to_vec()).collect();

    // Normalize vectors
    let norms: Vec<f32> = raw_vecs
        .iter()
        .map(|v| {
            let s: f32 = v.iter().map(|x| x * x).sum();
            s.sqrt().max(1e-10)
        })
        .collect();
    let normed: Vec<Vec<f32>> = raw_vecs
        .iter()
        .zip(norms.iter())
        .map(|(v, n)| v.iter().map(|x| x / n).collect())
        .collect();

    // Find similar groups
    let mut visited = vec![false; n];
    let mut groups: Vec<Vec<usize>> = Vec::new();

    for i in 0..n {
        if visited[i] { continue; }
        let mut group = vec![i];
        visited[i] = true;
        for j in (i + 1)..n {
            if visited[j] { continue; }
            let sim: f32 = normed[i].iter().zip(normed[j].iter()).map(|(a, b)| a * b).sum();
            if sim > threshold {
                group.push(j);
                visited[j] = true;
            }
        }
        if group.len() > 1 {
            groups.push(group);
        }
    }

    if groups.is_empty() {
        println!("\n\u{2713} No similar chunks found. Knowledge is already compact.");
        return Ok(());
    }

    let total_mergeable: usize = groups.iter().map(|g| g.len()).sum();
    let savings = total_mergeable - groups.len();
    println!("\nFound {} groups ({} chunks \u{2192} {} merged, saving {})",
        groups.len(), total_mergeable, groups.len(), savings);

    let llm = ChatClient::new(
        std::env::var("LLM_API_BASE").unwrap_or_else(|_| "https://api.riotpiao.com".to_string()),
        std::env::var("LLM_API_KEY").unwrap_or_default(),
        "reasoning",
    )?;

    let mut to_remove: Vec<usize> = Vec::new();

    for (gi, group) in groups.iter().enumerate() {
        let group_texts: Vec<&str> = group.iter().map(|&i| texts[i].as_str()).collect();
        let group_sources: Vec<&str> = group
            .iter()
            .map(|&i| records[i]["data"]["source"].as_str().unwrap_or("?"))
            .collect();

        // Compute max similarity in group
        let mut max_sim: f32 = 0.0;
        for a in 0..group.len() {
            for b in (a + 1)..group.len() {
                let sim: f32 = normed[group[a]].iter().zip(normed[group[b]].iter()).map(|(x, y)| x * y).sum();
                max_sim = max_sim.max(sim);
            }
        }

        println!("\nGroup {} (sim={:.3}, {} chunks):", gi + 1, max_sim, group.len());
        for &idx in group {
            let preview: String = texts[idx].chars().take(80).collect();
            let src = std::path::Path::new(group_sources[group.iter().position(|&i| i == idx).unwrap()])
                .file_stem()
                .map(|s| s.to_string_lossy().to_string())
                .unwrap_or_else(|| "?".to_string());
            println!("  [{}] {}...", src, preview.replace('\n', " "));
        }

        if dry_run {
            continue;
        }

        // Merge via LLM
        println!("  \u{2192} Merging with reasoning model...");
        let numbered: String = group_texts
            .iter()
            .zip(group_sources.iter())
            .enumerate()
            .map(|(i, (t, s))| format!("[Chunk {} from {}]:\n{}", i + 1, s, t))
            .collect::<Vec<_>>()
            .join("\n\n");

        let merged = llm.complete(
            "Merge these similar knowledge chunks into ONE concise chunk. Keep ALL unique facts. \
             Remove redundancy. Keep markdown formatting. Output ONLY the merged text.",
            &format!("Merge these {} chunks:\n\n{}", group.len(), numbered),
            1500,
        ).await?;

        // Strip <think> tags
        let merged_text = merged.text.split("</think>").last().unwrap_or(&merged.text).trim().to_string();
        let old_size: usize = group_texts.iter().map(|t| t.len()).sum();
        println!("  \u{2192} Merged: {} chars (was {} chars, {:.0}% reduction)",
            merged_text.len(), old_size, (1.0 - merged_text.len() as f64 / old_size as f64) * 100.0);

        // Update first chunk with merged content
        let mut hasher = Sha256::new();
        hasher.update(merged_text.as_bytes());
        let new_hash = format!("{:x}", hasher.finalize());

        records[group[0]]["data"]["text"] = serde_json::Value::String(merged_text);
        records[group[0]]["data"]["sha256"] = serde_json::Value::String(new_hash);
        records[group[0]]["data"]["merged_from"] = serde_json::json!(group.len());

        // Mark rest for removal
        for &idx in &group[1..] {
            to_remove.push(idx);
        }
    }

    if dry_run {
        println!("\n(dry run \u{2014} no changes written)");
        return Ok(());
    }

    // Write compacted log
    let to_remove_set: std::collections::HashSet<usize> = to_remove.into_iter().collect();
    let compacted: Vec<&serde_json::Value> = records
        .iter()
        .enumerate()
        .filter(|(i, _)| !to_remove_set.contains(i))
        .map(|(_, r)| r)
        .collect();

    // Backup
    let backup = format!("{}.bak", log_path);
    fs::copy(&log_path, &backup)?;

    // Write
    let mut f = fs::File::create(&log_path)?;
    use std::io::Write;
    for r in &compacted {
        serde_json::to_writer(&mut f, r)?;
        f.write_all(b"\n")?;
    }

    println!("\n{}", "\u{2500}".repeat(50));
    println!("Before: {} chunks", records.len());
    println!("After:  {} chunks (-{})", compacted.len(), records.len() - compacted.len());
    println!("Backup: {}", backup);
    println!("Written: {}", log_path);

    Ok(())
}

fn cmd_learn(
    paths: &[PathBuf],
    project: &str,
    dry_run: bool,
    max_chunk: usize,
) -> anyhow::Result<()> {
    use sha2::{Digest, Sha256};

    let mut all_files: Vec<PathBuf> = Vec::new();
    for p in paths {
        if p.is_dir() {
            for entry in walkdir::WalkDir::new(p)
                .into_iter()
                .filter_map(|e| e.ok())
                .filter(|e| {
                    e.path()
                        .extension()
                        .map(|ext| ext == "md")
                        .unwrap_or(false)
                })
            {
                all_files.push(entry.into_path());
            }
        } else if p.extension().map(|e| e == "md").unwrap_or(false) {
            all_files.push(p.clone());
        } else {
            eprintln!("Skipping non-markdown file: {}", p.display());
        }
    }

    if all_files.is_empty() {
        eprintln!("No markdown files found.");
        return Ok(());
    }

    all_files.sort();
    println!("Found {} markdown files", all_files.len());

    let mut total_chunks = 0usize;
    let mut total_bytes = 0usize;
    let mut log = if !dry_run {
        Some(mem_store::LogWriter::new(project, "learn", "latest")?)
    } else {
        None
    };

    for file in &all_files {
        let content = fs::read_to_string(file)?;
        let filename = file.file_stem().unwrap().to_string_lossy();
        let chunks = chunk_markdown(&content, max_chunk);

        println!("\n📄 {} — {} chunks", file.display(), chunks.len());

        for (i, chunk) in chunks.iter().enumerate() {
            let mut hasher = Sha256::new();
            hasher.update(chunk.as_bytes());
            let hash = format!("{:x}", hasher.finalize());
            let short_hash = &hash[..12];

            total_chunks += 1;
            total_bytes += chunk.len();

            if dry_run {
                let preview: String = chunk.chars().take(80).collect();
                println!(
                    "  [{}/{}] {} ({} bytes) {}",
                    i + 1,
                    chunks.len(),
                    short_hash,
                    chunk.len(),
                    preview.replace('\n', " ")
                );
            } else {
                let record = mem_store::EventRecord {
                    project: project.to_string(),
                    query: format!("{}:{}", filename, i),
                    run: "latest".to_string(),
                    turn: i as u32,
                    event_type: "learn".to_string(),
                    data: serde_json::json!({
                        "source": file.to_string_lossy(),
                        "chunk_index": i,
                        "total_chunks": chunks.len(),
                        "sha256": hash,
                        "level": "L1",
                        "text": chunk,
                    }),
                };
                log.as_mut().unwrap().log(record)?;
                println!("  ✓ [{}/{}] {} ({} bytes)", i + 1, chunks.len(), short_hash, chunk.len());
            }
        }
    }

    println!("\n{}", "─".repeat(50));
    println!(
        "{} files → {} chunks ({:.1} KB)",
        all_files.len(),
        total_chunks,
        total_bytes as f64 / 1024.0
    );
    if dry_run {
        println!("(dry run — nothing written)");
    } else {
        println!("Written to log/{}/learn/latest.jsonl", project);
    }
    Ok(())
}

/// Split markdown on ## headings, respecting max_chunk size.
fn chunk_markdown(content: &str, max_chunk: usize) -> Vec<String> {
    let mut chunks = Vec::new();
    let mut current = String::new();

    for line in content.lines() {
        // Split on ## headings (keep # title in first chunk)
        if line.starts_with("## ") && !current.is_empty() {
            let trimmed = current.trim().to_string();
            if !trimmed.is_empty() {
                chunks.push(trimmed);
            }
            current = String::new();
        }

        current.push_str(line);
        current.push('\n');

        // Hard split if chunk too large
        if current.len() > max_chunk {
            let trimmed = current.trim().to_string();
            if !trimmed.is_empty() {
                chunks.push(trimmed);
            }
            current = String::new();
        }
    }

    let trimmed = current.trim().to_string();
    if !trimmed.is_empty() {
        chunks.push(trimmed);
    }
    chunks
}

fn cmd_sig(tool: &str, file: Option<&PathBuf>) -> anyhow::Result<()> {
    use mem_core::lesson;
    use std::io::Read;

    // Read failure log from file or stdin
    let mut output = String::new();
    if let Some(file_path) = file {
        output = fs::read_to_string(file_path)?;
    } else {
        std::io::stdin().read_to_string(&mut output)?;
    }

    // Extract signature
    match lesson::extract(tool, &output) {
        Some(sig) => {
            println!("=== Failure Signature ===");
            println!("Tool: {}", sig.tool);
            println!("Rule: {}", sig.rule);
            println!("Hash (SHA256): {}", sig.sig_sha);
            println!("\n=== Raw Error ===");
            println!("{}", sig.raw);
            println!("\n=== Normalised Form ===");
            println!("{}", sig.normalised);
        }
        None => {
            eprintln!("Failed to extract signature for tool: {}", tool);
            std::process::exit(1);
        }
    }

    Ok(())
}