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

mod lessons_cmd;
mod http_server;
mod endpoints;

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,

    /// Start HTTP server
    Serve {
        #[arg(long, default_value = "8080")]
        port: u16,
        #[arg(long, default_value = "test-key")]
        api_key: String,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    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::Serve { port, api_key } => {
            http_server::start_server(port, api_key).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(())
}