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, /// 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, #[arg(long)] cwd: Option, }, /// 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, /// Infer the tool from this command line #[arg(long)] cmd: Option, /// Read the failure log from a file instead of stdin #[arg(long)] file: Option, #[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, 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(()) }