diff --git a/pkgs/atlas-search/Cargo.lock b/pkgs/atlas-search/Cargo.lock index af5f4d26a..b72e4831c 100644 --- a/pkgs/atlas-search/Cargo.lock +++ b/pkgs/atlas-search/Cargo.lock @@ -31,6 +31,56 @@ dependencies = [ "memchr", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "atlas-core" version = "0.1.0" @@ -55,6 +105,18 @@ dependencies = [ "tokenizers", ] +[[package]] +name = "atlas-search" +version = "0.1.0" +dependencies = [ + "atlas-core", + "atlas-embed", + "clap", + "rusqlite", + "serde_json", + "ureq", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -116,6 +178,52 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +[[package]] +name = "clap" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9c751b79415d4e559e3d1fcf128e09e720eb673a06d26cf6f392d37d75b66e0" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "clap_lex" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "compact_str" version = "0.9.1" @@ -352,6 +460,12 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "http" version = "1.5.0" @@ -384,6 +498,12 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.14.0" @@ -560,6 +680,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "onig" version = "6.5.3" @@ -785,7 +911,7 @@ dependencies = [ "getrandom 0.2.17", "libc", "untrusted", - "windows-sys", + "windows-sys 0.52.0", ] [[package]] @@ -1104,6 +1230,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "vcpkg" version = "0.2.15" @@ -1155,6 +1287,15 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-targets" version = "0.52.6" diff --git a/pkgs/atlas-search/Cargo.toml b/pkgs/atlas-search/Cargo.toml index 91775b88b..e85ec2786 100644 --- a/pkgs/atlas-search/Cargo.toml +++ b/pkgs/atlas-search/Cargo.toml @@ -4,7 +4,7 @@ # Self-contained workspace on the pkgs/id pattern. Crates land one step at # a time; step 1 (#5351) is atlas-embed only. [workspace] -members = ["crates/atlas-embed", "crates/atlas-core"] +members = ["crates/atlas-embed", "crates/atlas-core", "crates/atlas-search"] resolver = "2" [workspace.package] diff --git a/pkgs/atlas-search/crates/atlas-search/Cargo.toml b/pkgs/atlas-search/crates/atlas-search/Cargo.toml new file mode 100644 index 000000000..a438661e8 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "atlas-search" +version = "0.1.0" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description = "The Atlas search binary: one-shot CLI (search, get-chunk) at Python surface parity, standalone-first with a prefer-service mode; the serve subcommand lands in step 4. Epic developing-today/code#5350, step 3 (#5353)." + +[[bin]] +name = "atlas-search" +path = "src/main.rs" + +[dependencies] +atlas-core = { path = "../atlas-core" } +atlas-embed = { path = "../atlas-embed" } +clap = { version = "4", features = ["derive", "env"] } +serde_json = { version = "1", features = ["preserve_order"] } +ureq = { version = "3", features = ["rustls"] } + +[dev-dependencies] +rusqlite = { version = "0.37", features = ["bundled"] } +serde_json = { version = "1", features = ["preserve_order"] } diff --git a/pkgs/atlas-search/crates/atlas-search/src/main.rs b/pkgs/atlas-search/crates/atlas-search/src/main.rs new file mode 100644 index 000000000..029026933 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/src/main.rs @@ -0,0 +1,524 @@ +//! atlas-search — the Atlas access layer: hybrid search over the corpus. +//! +//! The one-shot CLI (epic developing-today/code#5350, step 3, +//! #5353): a surface-parity port of the Python `atlas_search.py`'s +//! front-end over atlas-core (step 2) and atlas-embed (step 1). +//! Standalone is the default and is complete — embed the query +//! in-process, talk to Turso or a local publish file directly. +//! `--via-service` prefers a running service and falls back to +//! standalone loudly until the step-4 protocol lands (#5354). +//! +//! Exit behavior mirrors the Python: `die()` fatals print +//! `atlas-search: ` on stderr and exit 2 (argparse errors +//! likewise exit 2 via clap); get-chunk's not-found exits 1, as do +//! backend/transport failures (Python tracebacks, also exit 1 — +//! the binary prints the message instead of the traceback). + +mod pyjson; +mod render; +mod service; + +use std::path::PathBuf; +use std::process::ExitCode; +use std::sync::OnceLock; + +use atlas_core::filters::{Filters, split_csv}; +use atlas_core::turso::{TURSO_HOST, resolve_turso_token}; +use atlas_core::{ + AtlasEmbedder, Error as CoreError, FixedEmbedder, LocalBackend, QueryEmbedder, SearchBackend, + SearchMode, SearchOptions, TursoBackend, VectorBackend, run_get_chunk, run_search, +}; +use clap::{Args, Parser, Subcommand, ValueEnum}; +use serde_json::Value as Json; + +// ---------------------------------------------------------------- args -- + +#[derive(Parser)] +#[command( + name = "atlas-search", + version, + about = "atlas-search — the Atlas access layer: hybrid search over the corpus." +)] +struct Cli { + #[command(subcommand)] + cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + /// hybrid FTS+vector search with classification filters + Search(Box), + /// one chunk with provenance + GetChunk(GetChunkArgs), +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +enum BackendChoice { + Auto, + Turso, + Local, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +enum ModeChoice { + Hybrid, + Fts, + Vector, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +enum EmbedderChoice { + Auto, + Fastembed, + St, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +enum VectorBackendChoice { + Scan, + IndexFill, + IndexFixed, +} + +/// Backend selection + via-service routing, shared by both +/// subcommands (the Python's `add_backend_args` + step 3's +/// service routing). +#[derive(Args, Clone)] +struct BackendArgs { + /// data backend (auto: turso when a token resolves, else local) + #[arg(long, value_enum, default_value_t = BackendChoice::Auto)] + backend: BackendChoice, + /// local backend: SQLite DB path (default: canonical atlas.db) + #[arg(long)] + db: Option, + /// local backend: comma-separated vector DB path(s) if the corpus DB has no vec table + #[arg(long)] + vecdb: Option, + /// prefer a running atlas service; fall back to standalone, loudly + #[arg(long)] + via_service: bool, + /// service base URL for --via-service (env ATLAS_SEARCH_URL) + #[arg(long)] + service_url: Option, +} + +/// The Python's `add_vector_args` (search only). +#[derive(Args, Clone)] +struct VectorArgs { + /// embedding model (default Snowflake/snowflake-arctic-embed-m-v2.0; env ATLAS_SEARCH_MODEL) + #[arg(long, env = "ATLAS_SEARCH_MODEL", default_value = atlas_core::DEFAULT_MODEL)] + model: String, + /// vector table name (default vec_snowflake_m_v2_768_f32; env ATLAS_VEC_TABLE) + #[arg(long, env = "ATLAS_VEC_TABLE", default_value = atlas_core::DEFAULT_VEC_TABLE)] + vec_table: String, + /// query embedder: fastembed (bge) or st = sentence-transformers with prompt_name='query' (Snowflake); auto picks by model name + #[arg(long, value_enum, default_value_t = EmbedderChoice::Auto)] + embedder: EmbedderChoice, + /// truncate query+doc vectors to N dims (Matryoshka; 0 = full) + #[arg(long, default_value_t = 0)] + trunc: i64, + /// turso vector branch: scan = exact vector_distance_cos over the filtered set; index-fixed = one vector_top_k lookup at pool depth, filtered afterwards; index-fill = vector_top_k with escalating depth until the pool fills after filtering (cap 3200) + #[arg(long, value_enum, default_value_t = VectorBackendChoice::Scan)] + vector_backend: VectorBackendChoice, +} + +#[derive(Args)] +struct SearchArgs { + query: String, + /// final results (default 10) + #[arg(long, default_value_t = 10)] + k: i64, + /// per-branch candidates (default 100) + #[arg(long, default_value_t = 100)] + pool: i64, + #[arg(long, value_enum, default_value_t = ModeChoice::Hybrid)] + mode: ModeChoice, + #[arg(long)] + kind: Option, + #[arg(long)] + topic: Option, + #[arg(long)] + project: Option, + #[arg(long)] + client: Option, + /// YYYY-MM-DD, on sessions.created_at (inclusive) + #[arg(long)] + since: Option, + /// YYYY-MM-DD, on sessions.created_at (inclusive) + #[arg(long)] + until: Option, + /// ceiling in shareable`. +type CliError = (u8, String); + +fn map_err(e: CoreError) -> CliError { + match e { + CoreError::Fatal(m) => (2, m), + CoreError::Backend(m) | CoreError::Embed(m) => (1, m), + } +} + +/// get-chunk's not-found is the Python's one `die(..., code=1)`. +fn map_get_chunk_err(e: CoreError, chunk_id: i64) -> CliError { + if let CoreError::Fatal(m) = &e + && *m == format!("chunk {chunk_id} not found") + { + return (1, m.clone()); + } + map_err(e) +} + +// ------------------------------------------------------------ backends -- + +fn home_dir() -> PathBuf { + std::env::var_os("HOME") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("~")) +} + +/// The Python's `default_db_path`. +fn default_db_path() -> String { + if let Ok(env) = std::env::var("ATLAS_DB") + && !env.is_empty() + { + return env; + } + let laptop = home_dir().join("atlas/turso-atlas/data/atlas.db"); + if laptop.exists() { + return laptop.to_string_lossy().into_owned(); + } + // The Python's script-relative ../data/atlas.db; the binary's + // analogue is its own directory's ../data. + if let Ok(exe) = std::env::current_exe() + && let Some(dir) = exe.parent().and_then(|p| p.parent()) + { + let cand = dir.join("data/atlas.db"); + if cand.exists() { + return cand.to_string_lossy().into_owned(); + } + } + laptop.to_string_lossy().into_owned() +} + +/// The Python's `pick_backend`: explicit turso/local, else turso +/// when a token resolves through the standard chain, else local. +fn pick_backend( + args: &BackendArgs, + vec_table: &str, + trunc: usize, +) -> Result, CliError> { + let local = |args: &BackendArgs| -> Result, CliError> { + let db = args.db.clone().unwrap_or_else(default_db_path); + LocalBackend::new(&db, args.vecdb.as_deref(), vec_table, trunc) + .map(|b| Box::new(b) as Box) + .map_err(map_err) + }; + match args.backend { + BackendChoice::Turso => TursoBackend::new(vec_table) + .map(|b| Box::new(b) as Box) + .map_err(map_err), + BackendChoice::Local => local(args), + BackendChoice::Auto => match resolve_turso_token() { + Ok((token, _source)) => Ok(Box::new(TursoBackend::with_token( + TURSO_HOST, token, vec_table, + ))), + Err(_) => local(args), + }, + } +} + +// ------------------------------------------------------------ embedder -- + +/// The pinned model revision (fixtures/README.md; step 5 fetches +/// it into the state dir — design section 4, Q1 as ruled). +const MODEL_REV: &str = "95c2741480856aa9666782eb4afe11959938017f"; + +fn resolve_model_dir() -> Option { + if let Ok(d) = std::env::var("ATLAS_MODEL_DIR") + && !d.is_empty() + { + let p = PathBuf::from(&d); + return if p.is_dir() { Some(p) } else { None }; + } + let home = home_dir(); + let state = home.join(format!(".local/state/atlas-search/model/{MODEL_REV}")); + if state.is_dir() { + return Some(state); + } + let hf = home.join(format!( + ".cache/huggingface/hub/models--Snowflake--snowflake-arctic-embed-m-v2.0/snapshots/{MODEL_REV}" + )); + if hf.is_dir() { + return Some(hf); + } + None +} + +fn resolve_ort_lib() -> Option { + if let Ok(p) = std::env::var("ATLAS_ORT_LIB") + && !p.is_empty() + { + let pb = PathBuf::from(&p); + return if pb.is_file() { Some(pb) } else { None }; + } + let cand = home_dir().join(".local/state/atlas-search/ort/libonnxruntime.so"); + if cand.is_file() { + return Some(cand); + } + None +} + +enum LoadedEmbedder { + Real(Box), + Fixed(FixedEmbedder), +} + +/// The CLI's lazy query embedder: loads the ONNX session on first +/// use only, so FTS-only searches and get-chunk never pay for (or +/// fail on) a missing model — the Python's degradation shape. +/// +/// Test hook: `ATLAS_SEARCH_QVEC_FILE` (a JSON array of floats) +/// substitutes a fixed query vector — the same monkeypatch the +/// Python fixture generator applies to `embed_query`; used by the +/// CLI parity tests, meaningless in production. +struct CliEmbedder { + cell: OnceLock>, + model: String, + embedder: EmbedderChoice, +} + +impl CliEmbedder { + fn new(model: String, embedder: EmbedderChoice) -> CliEmbedder { + CliEmbedder { + cell: OnceLock::new(), + model, + embedder, + } + } + + fn load(&self) -> Result<&LoadedEmbedder, CoreError> { + self.cell + .get_or_init(|| { + if let Ok(f) = std::env::var("ATLAS_SEARCH_QVEC_FILE") + && !f.is_empty() + { + let text = std::fs::read_to_string(&f) + .map_err(|e| format!("ATLAS_SEARCH_QVEC_FILE: cannot read {f}: {e}"))?; + let raw: Vec = serde_json::from_str::>(&text) + .map(|v| v.into_iter().map(|x| x as f32).collect()) + .map_err(|e| { + format!("ATLAS_SEARCH_QVEC_FILE: bad JSON vector in {f}: {e}") + })?; + return Ok(LoadedEmbedder::Fixed(FixedEmbedder { raw })); + } + load_embedder(&self.model, self.embedder).map(|e| LoadedEmbedder::Real(Box::new(e))) + }) + .as_ref() + .map_err(|m| CoreError::fatal(m.clone())) + } +} + +impl QueryEmbedder for CliEmbedder { + fn embed_query(&self, text: &str, trunc: usize) -> Result, CoreError> { + match self.load()? { + LoadedEmbedder::Real(e) => e.embed_query(text, trunc), + LoadedEmbedder::Fixed(e) => e.embed_query(text, trunc), + } + } +} + +/// Build the production embedder, with the Python's loud-failure +/// shape for the combinations it also refuses (an embedder it +/// cannot construct is a die(), never a silent downgrade). +fn load_embedder(model: &str, embedder: EmbedderChoice) -> Result { + if embedder == EmbedderChoice::Fastembed { + return Err( + "vector/hybrid modes with a fastembed embedder are not available in this \ + binary (the bge/fastembed fallback lives in the Python CLI); use --mode fts \ + or the Python atlas-search" + .to_string(), + ); + } + if model != atlas_core::DEFAULT_MODEL { + return Err(format!( + "vector/hybrid modes: this binary embeds only {} (atlas-embed); \ + model {model} needs the Python CLI", + atlas_core::DEFAULT_MODEL + )); + } + let model_dir = resolve_model_dir().ok_or_else(|| { + "Snowflake model not found (looked in ATLAS_MODEL_DIR, \ + ~/.local/state/atlas-search/model/, and the HF cache snapshot); \ + FTS-only search and get-chunk work without it" + .to_string() + })?; + let ort_lib = resolve_ort_lib().ok_or_else(|| { + "ONNX Runtime library not found: set ATLAS_ORT_LIB (the packaged \ + binary resolves it from its nix closure — step 5, #5355); \ + FTS-only search and get-chunk work without it" + .to_string() + })?; + let emb = atlas_embed::Embedder::load(&model_dir, &ort_lib) + .map_err(|e| format!("embedder load failed: {e}"))?; + Ok(AtlasEmbedder::new(emb)) +} + +// ------------------------------------------------------------- service -- + +/// When via-service is in play: probe, warn loudly on stderr, and +/// return the envelope note. Standalone proceeds regardless — +/// step 4 swaps the probe for the real protocol client. +fn via_service_note(args: &BackendArgs) -> Option { + if !(args.via_service || service::via_env()) { + return None; + } + let base = service::base_url(args.service_url.as_deref()); + let att = service::attempt(&base); + eprintln!( + "atlas-search: warning: --via-service: {}; falling back to standalone", + att.detail + ); + Some(format!("via-service: {}; served standalone", att.detail)) +} + +fn push_note(out: &mut Json, note: String) { + if let Some(obj) = out.as_object_mut() { + match obj.get_mut("notes") { + Some(Json::Array(notes)) => notes.push(Json::String(note)), + _ => { + obj.insert("notes".to_string(), Json::Array(vec![Json::String(note)])); + } + } + } +} + +// -------------------------------------------------------------- output -- + +fn print_output(out: &Json, text: bool) { + if text { + println!("{}", render::render_text(out)); + } else { + println!("{}", pyjson::dumps_pretty(out)); + } +} + +fn u(v: i64) -> usize { + usize::try_from(v).unwrap_or(0) +} + +// -------------------------------------------------------------- commands -- + +fn cmd_search(args: SearchArgs) -> Result<(), CliError> { + let note = via_service_note(&args.backend); + let filters = Filters::new( + split_csv(args.kind.as_deref()), + split_csv(args.topic.as_deref()), + split_csv(args.project.as_deref()), + split_csv(args.client.as_deref()), + &args.max_sensitivity, + args.since.clone(), + args.until.clone(), + ) + .map_err(map_err)?; + let trunc = u(args.vector.trunc); + let backend = pick_backend(&args.backend, &args.vector.vec_table, trunc)?; + let embedder = CliEmbedder::new(args.vector.model.clone(), args.vector.embedder); + let opts = SearchOptions { + query: args.query.clone(), + mode: match args.mode { + ModeChoice::Hybrid => SearchMode::Hybrid, + ModeChoice::Fts => SearchMode::Fts, + ModeChoice::Vector => SearchMode::Vector, + }, + k: u(args.k), + pool: u(args.pool), + group_by_session: args.group_by_session, + expand: u(args.expand), + with_message: args.with_message, + model: args.vector.model.clone(), + embedder_label: match args.vector.embedder { + EmbedderChoice::Auto => "auto", + EmbedderChoice::Fastembed => "fastembed", + EmbedderChoice::St => "st", + } + .to_string(), + trunc, + vector_backend: match args.vector.vector_backend { + VectorBackendChoice::Scan => VectorBackend::Scan, + VectorBackendChoice::IndexFill => VectorBackend::IndexFill, + VectorBackendChoice::IndexFixed => VectorBackend::IndexFixed, + }, + }; + let mut out = run_search(backend.as_ref(), &embedder, &filters, &opts).map_err(map_err)?; + if let Some(note) = note { + push_note(&mut out, note); + } + print_output(&out, args.text); + Ok(()) +} + +fn cmd_get_chunk(args: GetChunkArgs) -> Result<(), CliError> { + let note = via_service_note(&args.backend); + // The Python's get-chunk takes no vector args: pick_backend + // sees the defaults (getattr fallbacks in its main()). + let backend = pick_backend(&args.backend, atlas_core::DEFAULT_VEC_TABLE, 0)?; + let mut out = run_get_chunk(backend.as_ref(), args.chunk_id, u(args.neighbors)) + .map_err(|e| map_get_chunk_err(e, args.chunk_id))?; + if let Some(note) = note { + push_note(&mut out, note); + } + print_output(&out, args.text); + Ok(()) +} + +fn run(cli: Cli) -> Result<(), CliError> { + match cli.cmd { + Cmd::Search(args) => cmd_search(*args), + Cmd::GetChunk(args) => cmd_get_chunk(args), + } +} + +fn main() -> ExitCode { + let cli = ::parse(); + match run(cli) { + Ok(()) => ExitCode::SUCCESS, + Err((code, msg)) => { + eprintln!("atlas-search: {msg}"); + ExitCode::from(code) + } + } +} diff --git a/pkgs/atlas-search/crates/atlas-search/src/pyjson.rs b/pkgs/atlas-search/crates/atlas-search/src/pyjson.rs new file mode 100644 index 000000000..7a7c5a12d --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/src/pyjson.rs @@ -0,0 +1,266 @@ +//! Python-compatible JSON rendering for the CLI's output. +//! +//! The Python prints envelopes with +//! `json.dumps(out, indent=1, ensure_ascii=False)` and renders a +//! couple of `--text` fragments with the compact default dumps +//! (separators `", "` / `": "`, ensure_ascii=True). serde_json's +//! float formatting differs from CPython's repr in the small-value +//! range (serde prints `0.000087`, CPython prints `8.7e-05`), so +//! the writers here are hand-rolled over `serde_json::Value` with +//! CPython's float repr and escaping rules, making the binary's +//! stdout byte-identical to the Python's for the same envelope. + +use std::fmt::Write as _; + +use serde_json::Value as Json; + +/// CPython `repr(float)` (== `str(float)`): shortest round-trip +/// digits, decimal notation for decimal exponents in [-4, 16), +/// scientific outside that range with a signed, at-least-two-digit +/// exponent. +pub fn py_float(x: f64) -> String { + if x.is_nan() { + return "NaN".to_string(); + } + if x.is_infinite() { + return if x > 0.0 { + "Infinity".to_string() + } else { + "-Infinity".to_string() + }; + } + if x == 0.0 { + return if x.is_sign_negative() { + "-0.0".to_string() + } else { + "0.0".to_string() + }; + } + // Shortest digits + decimal exponent from Rust's scientific + // formatting (also shortest round-trip): "8.7e-5", "1.74242e-1". + let sci = format!("{:e}", x); + let (mant, exp_s) = sci.split_once('e').expect("LowerExp always has e"); + let exp: i32 = exp_s.parse().expect("exponent parses"); + let digits: String = mant.chars().filter(|c| c.is_ascii_digit()).collect(); + let mut s = String::new(); + if x < 0.0 { + s.push('-'); + } + if (-4..16).contains(&exp) { + if exp >= 0 { + let int_len = (exp + 1) as usize; + if digits.len() <= int_len { + s.push_str(&digits); + for _ in 0..(int_len - digits.len()) { + s.push('0'); + } + s.push_str(".0"); + } else { + s.push_str(&digits[..int_len]); + s.push('.'); + s.push_str(&digits[int_len..]); + } + } else { + s.push_str("0."); + for _ in 0..(-exp - 1) { + s.push('0'); + } + s.push_str(&digits); + } + } else { + s.push_str(&digits[..1]); + if digits.len() > 1 { + s.push('.'); + s.push_str(&digits[1..]); + } + s.push('e'); + s.push(if exp < 0 { '-' } else { '+' }); + let a = exp.abs(); + if a < 10 { + s.push('0'); + } + let _ = write!(s, "{a}"); + } + s +} + +/// A JSON string literal the way CPython's encoder writes it. +/// With `ensure_ascii`, non-ASCII is `\uXXXX`-escaped (surrogate +/// pairs above the BMP), as the compact dumps in render_text does; +/// without it, only the mandatory escapes apply (the envelope +/// dumps passes ensure_ascii=False). +pub fn py_string(s: &str, ensure_ascii: bool) -> String { + let mut out = String::with_capacity(s.len() + 2); + out.push('"'); + for c in s.chars() { + match c { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + '\u{08}' => out.push_str("\\b"), + '\u{0c}' => out.push_str("\\f"), + c if (c as u32) < 0x20 => { + let _ = write!(out, "\\u{:04x}", c as u32); + } + c if ensure_ascii && (c as u32) > 0x7f => { + if (c as u32) > 0xffff { + let v = (c as u32) - 0x10000; + let _ = write!( + out, + "\\u{:04x}\\u{:04x}", + 0xd800 + (v >> 10), + 0xdc00 + (v & 0x3ff) + ); + } else { + let _ = write!(out, "\\u{:04x}", c as u32); + } + } + c => out.push(c), + } + } + out.push('"'); + out +} + +fn write_number(out: &mut String, n: &serde_json::Number) { + if let Some(i) = n.as_i64() { + let _ = write!(out, "{i}"); + } else if let Some(u) = n.as_u64() { + let _ = write!(out, "{u}"); + } else if let Some(f) = n.as_f64() { + out.push_str(&py_float(f)); + } +} + +fn write_value(out: &mut String, v: &Json, pretty: bool, level: usize) { + match v { + Json::Null => out.push_str("null"), + Json::Bool(b) => out.push_str(if *b { "true" } else { "false" }), + Json::Number(n) => write_number(out, n), + Json::String(s) => out.push_str(&py_string(s, !pretty)), + Json::Array(items) => { + if items.is_empty() { + out.push_str("[]"); + return; + } + if pretty { + out.push_str("[\n"); + for (i, item) in items.iter().enumerate() { + if i > 0 { + out.push_str(",\n"); + } + out.push_str(&" ".repeat(level + 1)); + write_value(out, item, pretty, level + 1); + } + out.push('\n'); + out.push_str(&" ".repeat(level)); + out.push(']'); + } else { + out.push('['); + for (i, item) in items.iter().enumerate() { + if i > 0 { + out.push_str(", "); + } + write_value(out, item, pretty, level); + } + out.push(']'); + } + } + Json::Object(map) => { + if map.is_empty() { + out.push_str("{}"); + return; + } + if pretty { + out.push_str("{\n"); + for (i, (k, val)) in map.iter().enumerate() { + if i > 0 { + out.push_str(",\n"); + } + out.push_str(&" ".repeat(level + 1)); + out.push_str(&py_string(k, false)); + out.push_str(": "); + write_value(out, val, pretty, level + 1); + } + out.push('\n'); + out.push_str(&" ".repeat(level)); + out.push('}'); + } else { + out.push('{'); + for (i, (k, val)) in map.iter().enumerate() { + if i > 0 { + out.push_str(", "); + } + out.push_str(&py_string(k, true)); + out.push_str(": "); + write_value(out, val, pretty, level); + } + out.push('}'); + } + } + } +} + +/// `json.dumps(v, indent=1, ensure_ascii=False)` — the envelope +/// print (the caller adds print's trailing newline). +pub fn dumps_pretty(v: &Json) -> String { + let mut s = String::new(); + write_value(&mut s, v, true, 0); + s +} + +/// `json.dumps(v)` with default separators and ensure_ascii=True +/// — the `--text` header's filters fragment. +pub fn dumps_compact(v: &Json) -> String { + let mut s = String::new(); + write_value(&mut s, v, false, 0); + s +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn float_repr_matches_cpython() { + let cases = [ + (0.090909, "0.090909"), + (0.174242, "0.174242"), + (1.0, "1.0"), + (2.0, "2.0"), + (0.5, "0.5"), + (10.0, "10.0"), + (0.0001, "0.0001"), + (0.000087, "8.7e-05"), + (0.000001, "1e-06"), + (1e16, "1e+16"), + (1234567890123456.0, "1234567890123456.0"), + (0.3, "0.3"), + (0.0, "0.0"), + (-0.25, "-0.25"), + (100.0, "100.0"), + ]; + for (x, want) in cases { + assert_eq!(py_float(x), want, "py_float({x})"); + } + } + + #[test] + fn pretty_shape() { + let v = json!({"a": 1, "b": [1.5, "x"], "c": {}, "d": []}); + assert_eq!( + dumps_pretty(&v), + "{\n \"a\": 1,\n \"b\": [\n 1.5,\n \"x\"\n ],\n \"c\": {},\n \"d\": []\n}" + ); + } + + #[test] + fn compact_escapes_non_ascii() { + let v = json!({"k": "café"}); + assert_eq!(dumps_compact(&v), "{\"k\": \"caf\\u00e9\"}"); + assert_eq!(dumps_pretty(&v), "{\n \"k\": \"café\"\n}"); + } +} diff --git a/pkgs/atlas-search/crates/atlas-search/src/render.rs b/pkgs/atlas-search/crates/atlas-search/src/render.rs new file mode 100644 index 000000000..4afa50401 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/src/render.rs @@ -0,0 +1,187 @@ +//! The `--text` renderer: a line-for-line port of the Python's +//! `render_text` / `snippet`. Values render the way CPython's +//! f-strings render them (floats via repr, None as "None", +//! `rank_sources` as a Python dict repr). + +use serde_json::Value as Json; + +use crate::pyjson::{dumps_compact, py_float}; + +/// CPython f-string interpolation of a JSON value: strings raw, +/// null as `None`, bools as `True`/`False`, numbers via repr, +/// anything else via the compact dumps (the Python would repr a +/// dict/list; the fields rendered here are scalars in practice). +fn py_str(v: Option<&Json>) -> String { + match v { + None | Some(Json::Null) => "None".to_string(), + Some(Json::String(s)) => s.clone(), + Some(Json::Bool(b)) => { + if *b { + "True".to_string() + } else { + "False".to_string() + } + } + Some(Json::Number(n)) => { + if let Some(f) = n.as_f64() { + if n.is_i64() || n.is_u64() { + n.to_string() + } else { + py_float(f) + } + } else { + n.to_string() + } + } + Some(other) => dumps_compact(other), + } +} + +/// `{'fts': 1, 'vector': 2}` — a Python dict repr over the JSON +/// object, in the object's (insertion) order. +fn rank_sources_repr(v: Option<&Json>) -> String { + let Some(Json::Object(map)) = v else { + return "{}".to_string(); + }; + let parts: Vec = map + .iter() + .map(|(k, val)| format!("'{k}': {}", py_str(Some(val)))) + .collect(); + format!("{{{}}}", parts.join(", ")) +} + +/// The Python's `snippet`: whitespace-collapsed, 300 code points, +/// ellipsis when longer. +fn snippet(text: &str) -> String { + let collapsed = text.split_whitespace().collect::>().join(" "); + let mut chars = collapsed.chars(); + let head: String = chars.by_ref().take(300).collect(); + if chars.next().is_some() { + format!("{head}…") + } else { + head + } +} + +fn session_line(s: &Json, indent: &str) -> String { + let topics = match s.get("topics") { + Some(Json::Array(items)) => items + .iter() + .map(|t| match t { + Json::String(s) => s.clone(), + other => dumps_compact(other), + }) + .collect::>() + .join(","), + _ => String::new(), + }; + format!( + "{indent}{} · {} · {} · {} · {} · {}", + py_str(s.get("title")), + py_str(s.get("client")), + py_str(s.get("created_at")), + py_str(s.get("kind")), + topics, + py_str(s.get("sensitivity")), + ) +} + +/// The Python's `render_text` (which returns the joined lines; +/// the caller adds print's trailing newline). +pub fn render_text(out: &Json) -> String { + let mut lines: Vec = Vec::new(); + if out.get("query").is_some() { + lines.push(format!( + "query: {} mode: {} backend: {}", + py_str(out.get("query")), + py_str(out.get("mode")), + py_str(out.get("backend")), + )); + let empty = Json::Object(serde_json::Map::new()); + lines.push(format!( + "filters: {}", + dumps_compact(out.get("filters_applied").unwrap_or(&empty)) + )); + if let Some(notes) = out.get("notes").and_then(Json::as_array) { + for n in notes { + lines.push(format!("note: {}", py_str(Some(n)))); + } + } + if let Some(results) = out.get("results").and_then(Json::as_array) { + for (i, r) in results.iter().enumerate() { + let i = i + 1; + if r.get("chunks").is_some() { + // grouped by session + lines.push(format!( + "\n[{i}] session {} score={} hits={}", + py_str(r.get("session_id")), + py_str(r.get("score")), + py_str(r.get("hit_count")), + )); + if let Some(s) = r.get("session") { + lines.push(session_line(s, " ")); + } + if let Some(chunks) = r.get("chunks").and_then(Json::as_array) { + for ce in chunks { + lines.push(format!( + " chunk {} score={} {}", + py_str(ce.get("chunk_id")), + py_str(ce.get("score")), + rank_sources_repr(ce.get("rank_sources")), + )); + lines.push(format!( + " {}", + snippet(ce.get("text").and_then(Json::as_str).unwrap_or("")) + )); + } + } + } else { + lines.push(format!( + "\n[{i}] chunk {} score={} {}", + py_str(r.get("chunk_id")), + py_str(r.get("score")), + rank_sources_repr(r.get("rank_sources")), + )); + if let Some(s) = r.get("session") { + lines.push(session_line(s, " ")); + } + lines.push(format!( + " {}", + snippet(r.get("text").and_then(Json::as_str).unwrap_or("")) + )); + } + } + } + } else { + // get-chunk + lines.push(format!( + "chunk {} (idx {}) session {}", + py_str(out.get("chunk_id")), + py_str(out.get("chunk_idx")), + py_str(out.get("session_id")), + )); + if let Some(s) = out.get("session") { + lines.push(session_line(s, "")); + } + lines.push(format!( + "message {} role={}", + py_str(out.get("message").and_then(|m| m.get("id"))), + py_str(out.get("message").and_then(|m| m.get("role"))), + )); + lines.push(py_str(out.get("text"))); + if let Some(ctx) = out.get("context") { + for side in ["before", "after"] { + if let Some(items) = ctx.get(side).and_then(Json::as_array) { + for c in items { + lines.push(format!( + "[{side} chunk {}] {}", + py_str(c.get("chunk_id")), + snippet(c.get("text").and_then(Json::as_str).unwrap_or("")) + )); + } + } + } + } + } + lines.join("\n") +} diff --git a/pkgs/atlas-search/crates/atlas-search/src/service.rs b/pkgs/atlas-search/crates/atlas-search/src/service.rs new file mode 100644 index 000000000..029176759 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/src/service.rs @@ -0,0 +1,65 @@ +//! The via-service client seam (design section 1): *prefer +//! service, fall back to standalone loudly*. +//! +//! Step 3 (#5353) ships the contract, not the protocol: the job +//! protocol client proper lands in step 4 (#5354) against +//! atlas-proto, behind this same function. Today `attempt` probes +//! the configured base URL with a short timeout and reports what +//! it found; the caller then always serves standalone, printing +//! the warning on stderr and recording the note in the envelope +//! (matching how the Python reports backend fallbacks in +//! `notes`). No outcome of the probe — refused, timeout, an HTTP +//! answer from something that is not yet the service — can hang +//! or silently change a result. + +use std::time::Duration; + +/// What the probe found, in one clause, for the warning + note. +pub struct Attempt { + pub detail: String, +} + +/// The default service base URL (loopback bind per design +/// section 3; the port is proposed there pending the packaging +/// step's port-inventory check, #5355). +pub const DEFAULT_SERVICE_URL: &str = "http://127.0.0.1:18760"; + +/// Resolve the base URL: flag, then `ATLAS_SEARCH_URL`, then the +/// default. +pub fn base_url(flag: Option<&str>) -> String { + if let Some(u) = flag { + return u.trim_end_matches('/').to_string(); + } + if let Ok(u) = std::env::var("ATLAS_SEARCH_URL") + && !u.is_empty() + { + return u.trim_end_matches('/').to_string(); + } + DEFAULT_SERVICE_URL.to_string() +} + +/// Whether via-service was requested by env (`ATLAS_SEARCH_VIA=service`) +/// when the flag was not passed. +pub fn via_env() -> bool { + std::env::var("ATLAS_SEARCH_VIA").is_ok_and(|v| v == "service") +} + +/// Probe the service endpoint. Step 4 replaces this probe with +/// the real submit/page client; the fallback contract around it +/// does not change. +pub fn attempt(base: &str) -> Attempt { + let agent: ureq::Agent = ureq::Agent::config_builder() + .timeout_global(Some(Duration::from_secs(2))) + .timeout_connect(Some(Duration::from_secs(1))) + .http_status_as_error(false) + .build() + .into(); + let detail = match agent.get(base).call() { + Ok(resp) => format!( + "service at {base} answered HTTP {} but the job-protocol client lands in step 4 (#5354)", + resp.status().as_u16() + ), + Err(e) => format!("service at {base} unreachable ({e})"), + }; + Attempt { detail } +} diff --git a/pkgs/atlas-search/crates/atlas-search/tests/cli_parity.rs b/pkgs/atlas-search/crates/atlas-search/tests/cli_parity.rs new file mode 100644 index 000000000..648ed737f --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/tests/cli_parity.rs @@ -0,0 +1,237 @@ +//! CLI parity: the step-2 fixture corpus exercised through the +//! compiled `atlas-search` binary (developing-today/code#5353). +//! +//! For every search/get-chunk case recorded in +//! `fixtures/core-expectations.json` (envelopes produced by the +//! production Python), the binary runs against the fixture DB +//! with `--backend local` and the harness query-vector hook +//! (`ATLAS_SEARCH_QVEC_FILE` — the same monkeypatch the Python +//! fixture generator applies), and its stdout must parse to the +//! recorded envelope exactly. `--text` goldens (rendered by the +//! Python's own `render_text` from those envelopes) live in +//! `tests/golden/` and are compared byte-for-byte. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +use serde_json::Value as Json; + +fn fixtures_dir() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("../../fixtures") +} + +fn expectations() -> Json { + let p = fixtures_dir().join("core-expectations.json"); + serde_json::from_str(&std::fs::read_to_string(p).expect("read expectations")) + .expect("parse expectations") +} + +fn fixture_spec() -> Json { + let p = fixtures_dir().join("fixture_spec.json"); + serde_json::from_str(&std::fs::read_to_string(p).expect("read spec")).expect("parse spec") +} + +/// Build the fixture corpus DB in a unique temp dir (the same +/// builder the atlas-core parity tests use; the DB is never +/// committed). +fn build_fixture_db(tag: &str) -> PathBuf { + let dir = std::env::temp_dir().join(format!("atlas-cli-parity-{}-{tag}", std::process::id())); + std::fs::create_dir_all(&dir).expect("temp dir"); + let db = dir.join("fixture.db"); + let _ = std::fs::remove_file(&db); + let sql = std::fs::read_to_string(fixtures_dir().join("fixture_corpus.sql")) + .expect("read corpus sql"); + let spec = fixture_spec(); + let con = rusqlite::Connection::open(&db).expect("open fixture db"); + con.execute_batch(&sql).expect("load corpus sql"); + let vectors = spec["vectors"].as_object().expect("spec vectors"); + for (cid, vec) in vectors { + let mut blob = Vec::new(); + for x in vec.as_array().expect("vec array") { + blob.extend_from_slice(&(x.as_f64().expect("f64") as f32).to_le_bytes()); + } + con.execute( + "INSERT INTO vec_snowflake_m_v2_768_f32 (chunk_id, vec) VALUES (?, ?)", + rusqlite::params![cid.parse::().expect("cid"), blob], + ) + .expect("insert vec"); + } + drop(con); + db +} + +fn qvec_file() -> PathBuf { + let exp = expectations(); + let p = std::env::temp_dir().join(format!("atlas-cli-qvec-{}.json", std::process::id())); + std::fs::write(&p, serde_json::to_string(&exp["qvec_raw"]).unwrap()).expect("write qvec"); + p +} + +fn run_bin(db: &Path, args: &[&str]) -> std::process::Output { + let bin = env!("CARGO_BIN_EXE_atlas-search"); + Command::new(bin) + .args(args) + .arg("--db") + .arg(db) + .env("ATLAS_SEARCH_QVEC_FILE", qvec_file()) + .env_remove("ATLAS_VEC_DB") + .env_remove("ATLAS_SEARCH_VIA") + .env_remove("ATLAS_SEARCH_URL") + .env_remove("ATLAS_SEARCH_MODEL") + .env_remove("ATLAS_VEC_TABLE") + .output() + .expect("run atlas-search") +} + +/// Case name -> argv tail (after the subcommand), mirroring the +/// atlas-core parity test's `case_request`. +fn search_tail(name: &str) -> Vec<&'static str> { + match name { + "hybrid_default" => vec![], + "hybrid_secret_ceiling" => vec!["--max-sensitivity", "secret-bearing"], + "hybrid_shareable_ceiling" => vec!["--max-sensitivity", "shareable"], + "fts_default" => vec!["--mode", "fts"], + "vector_default" => vec!["--mode", "vector"], + "hybrid_kind_code" => vec!["--kind", "code"], + "hybrid_topic_nixos" => vec!["--topic", "nixos"], + "hybrid_client_codex_dates" => vec![ + "--client", + "codex", + "--since", + "2026-10-01", + "--until", + "2026-10-31", + ], + "hybrid_until_before_sB" => vec!["--client", "codex", "--until", "2026-10-19"], + "hybrid_grouped" => vec!["--group-by-session", "--k", "2"], + "hybrid_expand_message" => vec!["--expand", "1", "--with-message", "--k", "3"], + "hybrid_pool2_k2" => vec!["--pool", "2", "--k", "2"], + other => panic!("unknown case {other}"), + } +} + +#[test] +fn cli_search_envelopes_match_python() { + let db = build_fixture_db("search"); + let exp = expectations(); + for case in exp["search_cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let mut args = vec![ + "search", + "how do we back up the database", + "--backend", + "local", + ]; + args.extend(search_tail(name)); + let out = run_bin(&db, &args); + assert!( + out.status.success(), + "case {name}: exit {:?}, stderr: {}", + out.status.code(), + String::from_utf8_lossy(&out.stderr) + ); + let got: Json = serde_json::from_slice(&out.stdout) + .unwrap_or_else(|e| panic!("case {name}: stdout not JSON: {e}")); + assert_eq!( + got, + case["envelope"], + "case {name} envelope mismatch\ncli: {}\npython: {}", + serde_json::to_string_pretty(&got).unwrap(), + serde_json::to_string_pretty(&case["envelope"]).unwrap(), + ); + } +} + +#[test] +fn cli_get_chunk_matches_python() { + let db = build_fixture_db("getchunk"); + let exp = expectations(); + for case in exp["get_chunk_cases"].as_array().unwrap() { + let (cid, neighbors) = match case["name"].as_str().unwrap() { + "get_chunk_secret" => ("5", "0"), + "get_chunk_neighbors" => ("2", "1"), + other => panic!("unknown case {other}"), + }; + let out = run_bin( + &db, + &[ + "get-chunk", + cid, + "--neighbors", + neighbors, + "--backend", + "local", + ], + ); + assert!( + out.status.success(), + "case {}: exit {:?}, stderr: {}", + case["name"], + out.status.code(), + String::from_utf8_lossy(&out.stderr) + ); + let got: Json = serde_json::from_slice(&out.stdout).expect("stdout JSON"); + assert_eq!(got, case["envelope"], "case {}", case["name"]); + } +} + +#[test] +fn cli_get_chunk_not_found_exits_1() { + let db = build_fixture_db("notfound"); + let out = run_bin(&db, &["get-chunk", "999", "--backend", "local"]); + assert_eq!(out.status.code(), Some(1)); + assert_eq!( + String::from_utf8_lossy(&out.stderr).trim(), + "atlas-search: chunk 999 not found" + ); +} + +#[test] +fn cli_text_goldens_match_python() { + let db = build_fixture_db("text"); + let cases: [(&str, &str); 3] = [ + ("hybrid_default", "text_search_flat.txt"), + ("hybrid_grouped", "text_search_grouped.txt"), + ("get_chunk_neighbors", "text_get_chunk.txt"), + ]; + let golden_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/golden"); + for (name, file) in cases { + let golden = std::fs::read_to_string(golden_dir.join(file)) + .unwrap_or_else(|e| panic!("read golden {file}: {e}")); + let out = if name == "get_chunk_neighbors" { + run_bin( + &db, + &[ + "get-chunk", + "2", + "--neighbors", + "1", + "--backend", + "local", + "--text", + ], + ) + } else { + let mut args = vec![ + "search", + "how do we back up the database", + "--backend", + "local", + "--text", + ]; + args.extend(search_tail(name)); + run_bin(&db, &args) + }; + assert!( + out.status.success(), + "golden {name}: exit {:?}, stderr: {}", + out.status.code(), + String::from_utf8_lossy(&out.stderr) + ); + assert_eq!( + String::from_utf8_lossy(&out.stdout), + golden, + "golden {file} mismatch" + ); + } +} diff --git a/pkgs/atlas-search/crates/atlas-search/tests/golden/text_get_chunk.txt b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_get_chunk.txt new file mode 100644 index 000000000..8719a0d58 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_get_chunk.txt @@ -0,0 +1,5 @@ +chunk 2 (idx 1) session sA +NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal +message mA1 role=user +restic snapshots go to the turso bucket nightly +[before chunk 1] how do we back up the database with restic diff --git a/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_flat.txt b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_flat.txt new file mode 100644 index 000000000..96882cdc1 --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_flat.txt @@ -0,0 +1,26 @@ +query: how do we back up the database mode: hybrid backend: local +filters: {"max_sensitivity": "personal"} + +[1] chunk 1 score=0.181818 {'fts': 1, 'vector': 1} + NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal + how do we back up the database with restic + +[2] chunk 8 score=0.166667 {'fts': 2, 'vector': 2} + NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal + database backup verification steps + +[3] chunk 2 score=0.076923 {'vector': 3} + NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal + restic snapshots go to the turso bucket nightly + +[4] chunk 4 score=0.071429 {'vector': 4} + Rust refactor session · codex · 2026-10-20T09:12:00Z · code · rust · shareable + refactoring the rust parser for speed + +[5] chunk 3 score=0.066667 {'vector': 5} + NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal + the backup script uses restic and rclone + +[6] chunk 7 score=0.0625 {'vector': 6} + Rust refactor session · codex · 2026-10-20T09:12:00Z · code · rust · shareable + cargo test failures in the queue worker diff --git a/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_grouped.txt b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_grouped.txt new file mode 100644 index 000000000..295a6178f --- /dev/null +++ b/pkgs/atlas-search/crates/atlas-search/tests/golden/text_search_grouped.txt @@ -0,0 +1,20 @@ +query: how do we back up the database mode: hybrid backend: local +filters: {"max_sensitivity": "personal"} + +[1] session sA score=0.181818 hits=4 + NixOS backup chat · claude-code · 2026-10-05T14:23:00Z · chat · nixos,backup · personal + chunk 1 score=0.181818 {'fts': 1, 'vector': 1} + how do we back up the database with restic + chunk 8 score=0.166667 {'fts': 2, 'vector': 2} + database backup verification steps + chunk 2 score=0.076923 {'vector': 3} + restic snapshots go to the turso bucket nightly + chunk 3 score=0.066667 {'vector': 5} + the backup script uses restic and rclone + +[2] session sB score=0.071429 hits=2 + Rust refactor session · codex · 2026-10-20T09:12:00Z · code · rust · shareable + chunk 4 score=0.071429 {'vector': 4} + refactoring the rust parser for speed + chunk 7 score=0.0625 {'vector': 6} + cargo test failures in the queue worker