memvid_cli/commands/
sketch.rs

1//! Sketch track commands for fast candidate generation.
2//!
3//! Commands:
4//! - `sketch build`: Build sketches for all frames without existing sketches
5//! - `sketch info`: Show statistics about the sketch track
6
7use std::path::PathBuf;
8use std::time::Instant;
9
10use anyhow::{Context, Result};
11use clap::{Args, Subcommand};
12use memvid_core::{Memvid, SketchVariant};
13use serde_json::json;
14
15/// Arguments for the `sketch` command
16#[derive(Args)]
17pub struct SketchArgs {
18    #[command(subcommand)]
19    pub command: SketchCommand,
20}
21
22/// Sketch subcommands
23#[derive(Subcommand)]
24pub enum SketchCommand {
25    /// Build sketches for all frames that don't have one
26    Build(SketchBuildArgs),
27    /// Show sketch track statistics
28    Info(SketchInfoArgs),
29}
30
31/// Arguments for `sketch build`
32#[derive(Args)]
33pub struct SketchBuildArgs {
34    /// Path to the memory file
35    #[arg(value_name = "FILE")]
36    pub file: PathBuf,
37
38    /// Sketch variant: small (32 bytes), medium (64 bytes), large (96 bytes)
39    #[arg(long, default_value = "small")]
40    pub variant: SketchVariantArg,
41
42    /// Output as JSON
43    #[arg(long)]
44    pub json: bool,
45}
46
47/// Arguments for `sketch info`
48#[derive(Args)]
49pub struct SketchInfoArgs {
50    /// Path to the memory file
51    #[arg(value_name = "FILE")]
52    pub file: PathBuf,
53
54    /// Output as JSON
55    #[arg(long)]
56    pub json: bool,
57}
58
59/// Sketch variant argument
60#[derive(Clone, Debug)]
61pub struct SketchVariantArg(pub SketchVariant);
62
63impl Default for SketchVariantArg {
64    fn default() -> Self {
65        Self(SketchVariant::Small)
66    }
67}
68
69impl std::str::FromStr for SketchVariantArg {
70    type Err = String;
71
72    fn from_str(s: &str) -> Result<Self, Self::Err> {
73        match s.to_lowercase().as_str() {
74            "small" | "s" | "32" => Ok(Self(SketchVariant::Small)),
75            "medium" | "m" | "64" => Ok(Self(SketchVariant::Medium)),
76            "large" | "l" | "96" => Ok(Self(SketchVariant::Large)),
77            _ => Err(format!(
78                "Unknown variant '{}'. Use: small (32 bytes), medium (64 bytes), or large (96 bytes)",
79                s
80            )),
81        }
82    }
83}
84
85/// Handle the sketch command
86pub fn handle_sketch(args: SketchArgs) -> Result<()> {
87    match args.command {
88        SketchCommand::Build(build_args) => handle_sketch_build(build_args),
89        SketchCommand::Info(info_args) => handle_sketch_info(info_args),
90    }
91}
92
93/// Handle `sketch build`
94fn handle_sketch_build(args: SketchBuildArgs) -> Result<()> {
95    let start = Instant::now();
96
97    let mut mem = Memvid::open(&args.file)
98        .with_context(|| format!("Failed to open memory: {}", args.file.display()))?;
99
100    let frame_count = mem.frame_count();
101    let existing_sketches = mem.sketches().len();
102
103    // Build sketches for frames that don't have one
104    let new_sketches = mem.build_all_sketches(args.variant.0);
105
106    // Commit if we added any sketches
107    if new_sketches > 0 {
108        mem.commit()
109            .context("Failed to commit sketch track changes")?;
110    }
111
112    let elapsed = start.elapsed();
113    let stats = mem.sketch_stats();
114
115    if args.json {
116        let output = json!({
117            "file": args.file.display().to_string(),
118            "variant": format!("{:?}", args.variant.0),
119            "frames_total": frame_count,
120            "sketches_existing": existing_sketches,
121            "sketches_new": new_sketches,
122            "sketches_total": stats.entry_count,
123            "size_bytes": stats.size_bytes,
124            "elapsed_ms": elapsed.as_millis(),
125        });
126        println!("{}", serde_json::to_string_pretty(&output)?);
127    } else {
128        println!("Sketch build complete for {}", args.file.display());
129        println!();
130        println!("  Variant:          {:?}", args.variant.0);
131        println!("  Total frames:     {}", frame_count);
132        println!("  Existing sketches: {}", existing_sketches);
133        println!("  New sketches:     {}", new_sketches);
134        println!("  Total sketches:   {}", stats.entry_count);
135        println!("  Size on disk:     {}", format_bytes(stats.size_bytes));
136        println!("  Time:             {:.2}s", elapsed.as_secs_f64());
137
138        if new_sketches == 0 && existing_sketches == stats.entry_count as usize {
139            println!();
140            println!("All frames already have sketches.");
141        } else if new_sketches > 0 {
142            println!();
143            let rate = new_sketches as f64 / elapsed.as_secs_f64();
144            println!("Rate: {:.0} sketches/sec", rate);
145        }
146    }
147
148    Ok(())
149}
150
151/// Handle `sketch info`
152fn handle_sketch_info(args: SketchInfoArgs) -> Result<()> {
153    let mem = Memvid::open_read_only(&args.file)
154        .with_context(|| format!("Failed to open memory: {}", args.file.display()))?;
155
156    let stats = mem.sketch_stats();
157    let frame_count = mem.frame_count();
158    let coverage = if frame_count > 0 {
159        stats.entry_count as f64 / frame_count as f64 * 100.0
160    } else {
161        0.0
162    };
163
164    if args.json {
165        let output = json!({
166            "file": args.file.display().to_string(),
167            "enabled": !mem.sketches().is_empty(),
168            "variant": format!("{:?}", stats.variant),
169            "entry_count": stats.entry_count,
170            "frame_count": frame_count,
171            "coverage_percent": coverage,
172            "size_bytes": stats.size_bytes,
173            "short_text_count": stats.short_text_count,
174            "entry_size_bytes": stats.variant.entry_size(),
175        });
176        println!("{}", serde_json::to_string_pretty(&output)?);
177    } else {
178        println!("Sketch Track Info: {}", args.file.display());
179        println!();
180
181        if mem.sketches().is_empty() {
182            println!("  Status: No sketches built");
183            println!();
184            println!(
185                "  Run `memvid sketch build {}` to enable fast candidate generation.",
186                args.file.display()
187            );
188        } else {
189            println!("  Status:           Enabled");
190            println!(
191                "  Variant:          {:?} ({} bytes/entry)",
192                stats.variant,
193                stats.variant.entry_size()
194            );
195            println!("  Entries:          {}", stats.entry_count);
196            println!(
197                "  Coverage:         {:.1}% ({}/{})",
198                coverage, stats.entry_count, frame_count
199            );
200            println!("  Size:             {}", format_bytes(stats.size_bytes));
201            println!("  Short text:       {} entries", stats.short_text_count);
202            println!();
203            println!(
204                "Sketch-enabled queries will filter {} candidates before BM25 rerank.",
205                stats.entry_count
206            );
207        }
208    }
209
210    Ok(())
211}
212
213fn format_bytes(bytes: u64) -> String {
214    if bytes >= 1_073_741_824 {
215        format!("{:.1} GB", bytes as f64 / 1_073_741_824.0)
216    } else if bytes >= 1_048_576 {
217        format!("{:.1} MB", bytes as f64 / 1_048_576.0)
218    } else if bytes >= 1024 {
219        format!("{:.1} KB", bytes as f64 / 1024.0)
220    } else {
221        format!("{} bytes", bytes)
222    }
223}
memvid_cli/commands/sketch.rs

memvid_cli/commands/
sketch.rs
