From c89c9bec4c1e73cc648630bce3bc531d39be87d8 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 03:40:17 +0530 Subject: [PATCH 01/16] =?UTF-8?q?feat:=20rebrand=20embedded=20skill=20card?= =?UTF-8?q?s=20ponytail=E2=86=92flare=20+=20FLARE=5FOUTPUT=5FMODEL=20env-v?= =?UTF-8?q?ar?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/flare-code/src/skill-no-hallucination.md | 14 +++++++------- crates/flare-code/src/skill-playbook.md | 11 ++++++----- crates/flare-code/src/skill-review.md | 10 +++++----- crates/flare-code/src/skill.md | 4 ++-- crates/flare-output/src/llm.rs | 4 +++- 5 files changed, 23 insertions(+), 20 deletions(-) diff --git a/crates/flare-code/src/skill-no-hallucination.md b/crates/flare-code/src/skill-no-hallucination.md index 399c65ab..e0e49843 100644 --- a/crates/flare-code/src/skill-no-hallucination.md +++ b/crates/flare-code/src/skill-no-hallucination.md @@ -1,15 +1,15 @@ --- -name: ponytail-no-hallucination +name: flare-no-hallucination description: > - Reality-check companion to ponytail. Blocks invented APIs, deprecated + Reality-check companion to flare-code. Blocks invented APIs, deprecated methods, framework confusion, and undeclared variables — a minimal-looking line that calls a function which doesn't exist is not lazy, it's a bug with - extra confidence. Use whenever the user says "ponytail-no-hallucination", - "/ponytail-no-hallucination", "no hallucinations", "verify APIs", or "don't - invent functions". + extra confidence. Use whenever the user says "flare-no-hallucination", + "/flare-no-hallucination", "no hallucinations", "verify APIs", or "don't + invent functions" (legacy: "ponytail-no-hallucination" also works). --- -# Ponytail — No Hallucination Layer +# Flare — No Hallucination Layer The true lazy path is to use only what is provably there. A one-liner that calls a function which doesn't exist isn't minimal, it's a confident bug. @@ -32,4 +32,4 @@ version before using it. Trust but verify: read the file or docs first, prefer searching the codebase over assuming the API surface. -"stop ponytail-no-hallucination" or "normal mode" to revert. +"stop flare-no-hallucination" / "stop ponytail-no-hallucination" or "normal mode" to revert. diff --git a/crates/flare-code/src/skill-playbook.md b/crates/flare-code/src/skill-playbook.md index 1e4b4a20..0545598a 100644 --- a/crates/flare-code/src/skill-playbook.md +++ b/crates/flare-code/src/skill-playbook.md @@ -1,10 +1,11 @@ --- -name: ponytail-playbook +name: flare-playbook description: > TDD-aware project companion. Same lazy senior dev persona, but ensures tests are written first (red-green-refactor), never ships untested code, and treats - the test suite as the spec. Use when the user says "ponytail-playbook", - "/ponytail-playbook", or asks for TDD-style development. + the test suite as the spec. Use when the user says "flare-playbook", + "/flare-playbook", or asks for TDD-style development + (legacy: "ponytail-playbook" also works). --- You are a lazy senior developer on a TDD-aware project. The test suite is the @@ -19,7 +20,7 @@ spec — every behavior change starts with a failing test. 5. **Tests are documentation.** Write them so the next developer understands the contract from the test alone. -## The ladder (same as ponytail, with TDD priors) +## The ladder (same as flare-code, with TDD priors) 1. Does this need to exist at all? (YAGNI applies to tests too — don't test the framework, don't test getters.) @@ -38,4 +39,4 @@ After each change: 3. The refactored result (if different from step 2) 4. What was skipped, when to add it -"stop ponytail-playbook" or "normal mode" to revert to standard ponytail. +"stop flare-playbook" / "stop ponytail-playbook" or "normal mode" to revert to standard flare-code. diff --git a/crates/flare-code/src/skill-review.md b/crates/flare-code/src/skill-review.md index 6325269e..fa6a2e64 100644 --- a/crates/flare-code/src/skill-review.md +++ b/crates/flare-code/src/skill-review.md @@ -1,13 +1,13 @@ --- -name: ponytail-review +name: flare-review description: > Code review focused exclusively on over-engineering. Finds what to delete: reinvented standard library, unneeded dependencies, speculative abstractions, dead flexibility. One line per finding: location, what to cut, what replaces it. Use when the user says "review for over-engineering", "what can we delete", "is this over-engineered", "simplify review", or invokes - /ponytail-review. Complements correctness-focused review, this one only - hunts complexity. + /flare-review. Complements correctness-focused review, this one only + hunts complexity (legacy: "ponytail-review" also works). --- Review diffs for unnecessary complexity. Number each finding sequentially. @@ -86,6 +86,6 @@ If there is nothing to cut, say `Lean already. Ship.` and stop. Scope: over-engineering and complexity only. Correctness bugs, security holes, and performance are explicitly out of scope. Route them to a normal review pass, not this one. A single smoke test or `assert`-based -self-check is the ponytail minimum, not bloat, never flag it for deletion. +self-check is the flare-code minimum, not bloat, never flag it for deletion. Does not apply the fixes, only lists them. -"stop ponytail-review" or "normal mode": revert to verbose review style. +"stop flare-review" / "stop ponytail-review" or "normal mode": revert to verbose review style. diff --git a/crates/flare-code/src/skill.md b/crates/flare-code/src/skill.md index 20dd4e1e..62e78fb7 100644 --- a/crates/flare-code/src/skill.md +++ b/crates/flare-code/src/skill.md @@ -27,7 +27,7 @@ code is the code never written. ACTIVE EVERY RESPONSE. No drift back to over-building. Still active if unsure. Off only: "stop flare code" / "normal mode". Default: **full**. -Switch: \`/flare code lite|full|ultra\`. +Switch: \`/flare-code lite|full|ultra\`. ## The ladder @@ -134,7 +134,7 @@ test, YAGNI applies to tests too. ## Boundaries -Flare Code governs what you build, not how you talk (pair with Caveman for +Flare Code governs what you build, not how you talk (pair with flare-output for terse prose). "stop flare code" / "normal mode": revert. Level persists until changed or session end. diff --git a/crates/flare-output/src/llm.rs b/crates/flare-output/src/llm.rs index ccfcbd2d..a5013d11 100644 --- a/crates/flare-output/src/llm.rs +++ b/crates/flare-output/src/llm.rs @@ -26,7 +26,9 @@ impl Llm for RealLlm { } fn call_via_api(api_key: &str, prompt: &str) -> Result { - let model = std::env::var("CAVEMAN_MODEL").unwrap_or_else(|_| "claude-sonnet-4-5".to_string()); + let model = std::env::var("FLARE_OUTPUT_MODEL") + .or_else(|_| std::env::var("CAVEMAN_MODEL")) + .unwrap_or_else(|_| "claude-sonnet-4-5".to_string()); let body = serde_json::json!({ "model": model, "max_tokens": 8192, From e133834ec06154cc57aca5917e245e217f049f89 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 03:45:29 +0530 Subject: [PATCH 02/16] =?UTF-8?q?chore:=20uninstall=20caveman=20=E2=80=94?= =?UTF-8?q?=20remove=20from=20KNOWN=5FCOMPRESSION=5FPLUGINS,=20update=20ba?= =?UTF-8?q?ckup=20path=20and=20prompt?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 4 ++++ crates/flare-code/src/config.rs | 17 ++++------------- crates/flare-code/src/instructions.rs | 24 +++++------------------- crates/flare-output/src/compress.rs | 2 +- crates/flare-output/src/prompt.rs | 4 ++-- 5 files changed, 16 insertions(+), 35 deletions(-) diff --git a/.gitignore b/.gitignore index 48ba12b2..3064d977 100644 --- a/.gitignore +++ b/.gitignore @@ -1,11 +1,15 @@ /target /.docs/ /domain-research/ +/skills/ /.worktrees/ # Machine-local docs output / notes, not part of the project /docs/ +# Competitive audit clones and reports — large, transient, per-machine +/.competitive-audit/ + # Local opencode config — carries per-machine MCP server paths, not shared /opencode.json diff --git a/crates/flare-code/src/config.rs b/crates/flare-code/src/config.rs index 58ce63cd..60e3fc15 100644 --- a/crates/flare-code/src/config.rs +++ b/crates/flare-code/src/config.rs @@ -51,9 +51,9 @@ pub fn normalize_extended_mode(mode: &str) -> Option { pub static ENV_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); /// Compression/persona plugins known to conflict with flare-code's own style -/// guidance if both are active — e.g. caveman's terse-prose rules vs +/// guidance if both are active — e.g. another tool's terse-prose rules vs /// flare-code's own output-shape rules. -const KNOWN_COMPRESSION_PLUGINS: &[&str] = &["caveman"]; +const KNOWN_COMPRESSION_PLUGINS: &[&str] = &[]; fn claude_dir() -> PathBuf { std::env::var("CLAUDE_CONFIG_DIR").map_or_else( @@ -215,17 +215,8 @@ mod tests { } #[test] - fn detects_caveman_in_settings_json() { - let _guard = ENV_TEST_LOCK - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let dir = std::env::temp_dir().join("flare-code_test_compression_conflict"); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write(dir.join("settings.json"), r#"{"plugins": ["caveman"]}"#).unwrap(); - unsafe { std::env::set_var("CLAUDE_CONFIG_DIR", &dir) }; - assert_eq!(detect_compression_plugins(), vec!["caveman"]); - unsafe { std::env::remove_var("CLAUDE_CONFIG_DIR") }; - std::fs::remove_dir_all(&dir).ok(); + fn detects_no_plugins_when_list_empty() { + assert!(detect_compression_plugins().is_empty()); } #[test] diff --git a/crates/flare-code/src/instructions.rs b/crates/flare-code/src/instructions.rs index 24406327..3e52bd1d 100644 --- a/crates/flare-code/src/instructions.rs +++ b/crates/flare-code/src/instructions.rs @@ -65,10 +65,9 @@ pub fn build(mode: &str, skill_path: Option<&Path>) -> Instructions { } } -/// If a known compression/persona plugin (e.g. caveman) is also wired into -/// the agent's settings, add a short note so the two don't read as -/// contradictory: flare-code governs code structure, the peer plugin governs -/// output style. +/// If a known compression/persona plugin is also wired into the agent's +/// settings, add a short note so the two don't read as contradictory: +/// flare-code governs code structure, the peer plugin governs output style. fn compression_deconfliction() -> String { let peers = config::detect_compression_plugins(); if peers.is_empty() { @@ -124,22 +123,9 @@ mod tests { } #[test] - #[allow(unsafe_code)] - fn build_appends_deconfliction_when_compression_plugin_present() { - let _guard = config::ENV_TEST_LOCK - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let dir = std::env::temp_dir().join("flare-code_test_instructions_compression"); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write(dir.join("settings.json"), r#"{"plugins": ["caveman"]}"#).unwrap(); - unsafe { std::env::set_var("CLAUDE_CONFIG_DIR", &dir) }; - + fn build_has_no_deconfliction_when_known_plugins_empty() { let ins = build("full", None); - assert!(ins.body.contains("Compression plugin coexistence")); - assert!(ins.body.contains("caveman")); - - unsafe { std::env::remove_var("CLAUDE_CONFIG_DIR") }; - std::fs::remove_dir_all(&dir).ok(); + assert!(!ins.body.contains("Compression plugin coexistence")); } #[test] diff --git a/crates/flare-output/src/compress.rs b/crates/flare-output/src/compress.rs index d97381d7..93dd3907 100644 --- a/crates/flare-output/src/compress.rs +++ b/crates/flare-output/src/compress.rs @@ -144,7 +144,7 @@ fn backup_path_for(target: &Path, mode: BackupMode) -> PathBuf { let base = dirs::cache_dir() .unwrap_or_else(|| PathBuf::from(".")) .join("agentflare") - .join("caveman") + .join("flare-output") .join("backups"); // Hash the full parent path, not just its last component — two // files with the same name under differently-located but diff --git a/crates/flare-output/src/prompt.rs b/crates/flare-output/src/prompt.rs index 855d6e9e..41381f70 100644 --- a/crates/flare-output/src/prompt.rs +++ b/crates/flare-output/src/prompt.rs @@ -1,4 +1,4 @@ -//! Compression prompt profiles. `Generic` is caveman's own embedded prompt +//! Compression prompt profiles. `Generic` is flare-output's embedded prompt //! — a short string literal that has never had an external doc, so there's //! no reason to invent a download mechanism for it (unlike ponytail's //! SKILL.md, which lives in its own separately-maintained upstream repo). @@ -28,7 +28,7 @@ impl Prompt { match self { Prompt::Generic => { format!( - "Compress this markdown into caveman format.\n\n{GENERIC_RULES}\n\nTEXT:\n{body}" + "Compress this markdown into flare format.\n\n{GENERIC_RULES}\n\nTEXT:\n{body}" ) } Prompt::Custom(spec) => format!( From bc72d0d8ea38a0c1fc228bf4ee2b781514b05b8e Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 13:34:32 +0530 Subject: [PATCH 03/16] feat(agentflare-store): initial crate - kv, documents (CRUD+FTS+vector), blobs, leases Imports the agentflare-store implementation for item #147, previously left as uncommitted working-tree state with zero git history. Store skeleton, kv, documents CRUD/FTS5/vector search/hybrid fusion, blob storage, and a leases re-export of db_kit::claim. Wires the crate into the workspace members list. Applies three plan-mandated fixes during import: drop deps unused anywhere in the crate (agentflare-artifacts, sha2, hex, uuid, chrono, mime_guess), use BLAKE3 instead of SHA-256 for blob hashing per the plan's global constraints, and a clippy nit in the hybrid-search chain (redundant into_iter()). Verified: cargo build --workspace --all-features, cargo clippy -p agentflare-store -- -D warnings, cargo test --workspace (27/27 in this crate, full workspace suite green). Known gaps vs the 8-task plan, tracked as follow-up work: - embeddings module is a 50-line stub (cosine_similarity/normalize only) - Task 5's vendoring of lean-ctx's model/tokenizer/download pipeline never happened, so doc_vec_search/doc_hybrid_search are untestable end-to-end - blob storage chunks bytes into SQLite rather than content-addressed files on disk under a blobs dir, as the plan specified - documents schema is missing title/doc_type/blob_hash/mime/tags/session_id/ source/version/history and has no versioning API - Store holds a bare rusqlite::Connection (not parking_lot::Mutex-wrapped), so it isn't Sync and can't be shared across threads as the plan's Arc design requires - Task 8 (migrate state.json onto this crate's kv store) hasn't been started; src/state.rs is untouched and agentflare-store isn't yet a dependency of the main package --- Cargo.lock | 115 ++++++ Cargo.toml | 2 +- crates/agentflare-store/Cargo.toml | 24 ++ crates/agentflare-store/src/blobs.rs | 157 ++++++++ crates/agentflare-store/src/documents.rs | 458 ++++++++++++++++++++++ crates/agentflare-store/src/embed.rs | 50 +++ crates/agentflare-store/src/kv.rs | 130 ++++++ crates/agentflare-store/src/leases.rs | 55 +++ crates/agentflare-store/src/lib.rs | 74 ++++ crates/agentflare-store/src/migrate.rs | 95 +++++ crates/agentflare-store/src/migrations.rs | 61 +++ 11 files changed, 1220 insertions(+), 1 deletion(-) create mode 100644 crates/agentflare-store/Cargo.toml create mode 100644 crates/agentflare-store/src/blobs.rs create mode 100644 crates/agentflare-store/src/documents.rs create mode 100644 crates/agentflare-store/src/embed.rs create mode 100644 crates/agentflare-store/src/kv.rs create mode 100644 crates/agentflare-store/src/leases.rs create mode 100644 crates/agentflare-store/src/lib.rs create mode 100644 crates/agentflare-store/src/migrate.rs create mode 100644 crates/agentflare-store/src/migrations.rs diff --git a/Cargo.lock b/Cargo.lock index 22d235ca..e6311a19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -221,6 +221,22 @@ dependencies = [ "thiserror", ] +[[package]] +name = "agentflare-store" +version = "0.1.0" +dependencies = [ + "agentflare-db-kit", + "blake3", + "bytemuck", + "parking_lot", + "rusqlite", + "rusqlite_migration", + "serde", + "serde_json", + "tempfile", + "thiserror", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -298,6 +314,18 @@ dependencies = [ "derive_arbitrary", ] +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.89" @@ -400,6 +428,20 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -424,6 +466,26 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "bytes" version = "1.12.0" @@ -599,6 +661,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -1408,6 +1476,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1541,6 +1618,29 @@ version = "4.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d211803b9b6b570f68772237e415a029d5a50c65d382910b879fb19d3271f94d" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "password-hash" version = "0.5.0" @@ -1719,6 +1819,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_users" version = "0.5.2" @@ -2034,6 +2143,12 @@ dependencies = [ "syn", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "serde" version = "1.0.228" diff --git a/Cargo.toml b/Cargo.toml index 0e91839e..0ee28420 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["crates/flare-code", "crates/agent-registry", "crates/skill-registry", "crates/gateway-registry", "crates/flare-output", "crates/agentflare-artifacts", "crates/agentflare-backend", "crates/agentflare-db-kit", "crates/flare-search-kit"] +members = ["crates/flare-code", "crates/agent-registry", "crates/skill-registry", "crates/gateway-registry", "crates/flare-output", "crates/agentflare-artifacts", "crates/agentflare-backend", "crates/agentflare-db-kit", "crates/flare-search-kit", "crates/agentflare-store"] resolver = "2" [package] diff --git a/crates/agentflare-store/Cargo.toml b/crates/agentflare-store/Cargo.toml new file mode 100644 index 00000000..f4ec889e --- /dev/null +++ b/crates/agentflare-store/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "agentflare-store" +version = "0.1.0" +edition = "2024" +rust-version = "1.91" +license = "MIT" +description = "Local-first document, artifact, blob, and memory storage for AI agent CLIs. SQLite index + content-addressed files. Zero external infra." +publish = false + +[dependencies] +db_kit = { package = "agentflare-db-kit", path = "../agentflare-db-kit" } +rusqlite = { version = "0.40", features = ["bundled"] } +rusqlite_migration = "2" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +blake3 = "1.5" +thiserror = "2" +parking_lot = "0.12" +bytemuck = { version = "1", features = ["derive"] } +[features] +default = [] + +[dev-dependencies] +tempfile = "3" diff --git a/crates/agentflare-store/src/blobs.rs b/crates/agentflare-store/src/blobs.rs new file mode 100644 index 00000000..889fc156 --- /dev/null +++ b/crates/agentflare-store/src/blobs.rs @@ -0,0 +1,157 @@ +use crate::Store; +use rusqlite::{OptionalExtension, params}; +use serde::Serialize; + +#[derive(Debug, Clone, Serialize)] +pub struct BlobMeta { + pub hash: String, + pub size: i64, + pub ref_count: i32, + pub created_at: i64, +} + +const CHUNK_SIZE: usize = 64 * 1024; // 64 KiB + +impl Store { + pub fn blob_store(&self, data: &[u8]) -> rusqlite::Result { + let hash = blake3::hash(data).to_hex().to_string(); + let now = db_kit::ids::now(); + + // Bump ref count if exists + let exists = self + .conn + .query_row( + "SELECT 1 FROM store_blobs WHERE hash = ?1", + params![hash], + |_| Ok(()), + ) + .optional()? + .is_some(); + + if exists { + self.conn.execute( + "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", + params![hash], + )?; + return Ok(hash); + } + + self.conn.execute( + "INSERT INTO store_blobs (hash, size, ref_count, created_at) VALUES (?1, ?2, 1, ?3)", + params![hash, data.len() as i64, now], + )?; + + for (i, chunk) in data.chunks(CHUNK_SIZE).enumerate() { + self.conn.execute( + "INSERT INTO store_blob_chunks (hash, chunk_index, data) VALUES (?1, ?2, ?3)", + params![hash, i as i64, chunk], + )?; + } + + Ok(hash) + } + + pub fn blob_get(&self, hash: &str) -> rusqlite::Result>> { + let meta: BlobMeta = match self + .conn + .query_row( + "SELECT hash, size, ref_count, created_at FROM store_blobs WHERE hash = ?1", + params![hash], + |row| { + Ok(BlobMeta { + hash: row.get(0)?, + size: row.get(1)?, + ref_count: row.get(2)?, + created_at: row.get(3)?, + }) + }, + ) + .optional()? + { + Some(m) => m, + None => return Ok(None), + }; + + let mut stmt = self + .conn + .prepare("SELECT data FROM store_blob_chunks WHERE hash = ?1 ORDER BY chunk_index")?; + let chunks: rusqlite::Result>> = + stmt.query_map(params![hash], |row| row.get(0))?.collect(); + + let mut buf = Vec::with_capacity(meta.size as usize); + for chunk in chunks? { + buf.extend_from_slice(&chunk); + } + Ok(Some(buf)) + } + + pub fn blob_ref(&self, hash: &str) -> rusqlite::Result { + let n = self.conn.execute( + "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", + params![hash], + )?; + Ok(n > 0) + } + + pub fn blob_unref(&self, hash: &str) -> rusqlite::Result { + let n = self.conn.execute( + "UPDATE store_blobs SET ref_count = ref_count - 1 WHERE hash = ?1 AND ref_count > 0", + params![hash], + )?; + if n > 0 { + self.conn.execute( + "DELETE FROM store_blobs WHERE hash = ?1 AND ref_count <= 0", + params![hash], + )?; + self.conn.execute( + "DELETE FROM store_blob_chunks WHERE hash = ?1", + params![hash], + )?; + } + Ok(n > 0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn store() -> Store { + Store::open_memory().unwrap() + } + + #[test] + fn store_and_retrieve() { + let s = store(); + let data = b"hello blob store"; + let hash = s.blob_store(data).unwrap(); + assert_eq!(hash.len(), 64); + + let retrieved = s.blob_get(&hash).unwrap().unwrap(); + assert_eq!(retrieved, data); + } + + #[test] + fn get_missing() { + let s = store(); + assert!(s.blob_get("0000").unwrap().is_none()); + } + + #[test] + fn dedup_same_content() { + let s = store(); + let h1 = s.blob_store(b"same").unwrap(); + let h2 = s.blob_store(b"same").unwrap(); + assert_eq!(h1, h2); + } + + #[test] + fn ref_unref() { + let s = store(); + let h = s.blob_store(b"data").unwrap(); + assert!(s.blob_ref(&h).unwrap()); + assert!(s.blob_unref(&h).unwrap()); + assert!(s.blob_unref(&h).unwrap()); + assert!(s.blob_get(&h).unwrap().is_none()); + } +} diff --git a/crates/agentflare-store/src/documents.rs b/crates/agentflare-store/src/documents.rs new file mode 100644 index 00000000..e5ef871f --- /dev/null +++ b/crates/agentflare-store/src/documents.rs @@ -0,0 +1,458 @@ +use crate::Store; +use rusqlite::{OptionalExtension, params}; +use serde::Serialize; + +#[derive(Debug, Clone, Serialize)] +pub struct Document { + pub id: String, + pub project_id: String, + pub path: String, + pub content: String, + pub created_at: i64, + pub updated_at: i64, + pub deleted_at: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct DocMatch { + pub id: String, + pub project_id: String, + pub path: String, + pub snippet: String, + pub score: f64, +} + +impl Store { + fn doc_sync_fts(&self, row_id: i64, content: &str) -> rusqlite::Result<()> { + // FTS5 has no REPLACE/UPSERT — delete any existing rowid first (no-op if fresh) + self.conn.execute( + "DELETE FROM store_docs_fts WHERE rowid = ?1", + params![row_id], + )?; + self.conn.execute( + "INSERT INTO store_docs_fts(rowid, content) VALUES (?1, ?2)", + params![row_id, content], + )?; + Ok(()) + } + pub fn doc_upsert( + &self, + project_id: &str, + path: &str, + content: &str, + ) -> rusqlite::Result { + let now = db_kit::ids::now(); + let id = db_kit::ids::new_id(); + + // Try to find existing by (project_id, path), else insert fresh + let existing = self + .conn + .query_row( + "SELECT id, rowid FROM store_documents WHERE project_id = ?1 AND path = ?2", + params![project_id, path], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + ) + .optional()?; + + if let Some((existing_id, rowid)) = existing { + self.conn.execute( + "UPDATE store_documents SET content = ?1, updated_at = ?2, deleted_at = NULL WHERE id = ?3", + params![content, now, existing_id], + )?; + self.doc_sync_fts(rowid, content)?; + self.doc_get(&existing_id).map(|o| o.unwrap()) + } else { + self.conn.execute( + "INSERT INTO store_documents (id, project_id, path, content, created_at, updated_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?5)", + params![id, project_id, path, content, now], + )?; + let rowid = self.conn.last_insert_rowid(); + self.doc_sync_fts(rowid, content)?; + Ok(Document { + id, + project_id: project_id.to_string(), + path: path.to_string(), + content: content.to_string(), + created_at: now, + updated_at: now, + deleted_at: None, + }) + } + } + + pub fn doc_get(&self, id: &str) -> rusqlite::Result> { + self.conn + .query_row( + "SELECT id, project_id, path, content, created_at, updated_at, deleted_at + FROM store_documents WHERE id = ?1", + params![id], + |row| { + Ok(Document { + id: row.get(0)?, + project_id: row.get(1)?, + path: row.get(2)?, + content: row.get(3)?, + created_at: row.get(4)?, + updated_at: row.get(5)?, + deleted_at: row.get(6)?, + }) + }, + ) + .optional() + } + + pub fn doc_delete(&self, id: &str) -> rusqlite::Result { + let now = db_kit::ids::now(); + if let Some(rowid) = self + .conn + .query_row( + "SELECT rowid FROM store_documents WHERE id = ?1", + params![id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + self.conn.execute( + "UPDATE store_documents SET deleted_at = ?1 WHERE id = ?2", + params![now, id], + )?; + self.conn.execute( + "DELETE FROM store_docs_fts WHERE rowid = ?1", + params![rowid], + )?; + Ok(true) + } else { + Ok(false) + } + } + + pub fn doc_hard_delete(&self, id: &str) -> rusqlite::Result { + if let Some(rowid) = self + .conn + .query_row( + "SELECT rowid FROM store_documents WHERE id = ?1", + params![id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + self.conn + .execute("DELETE FROM store_documents WHERE id = ?1", params![id])?; + self.conn.execute( + "DELETE FROM store_docs_fts WHERE rowid = ?1", + params![rowid], + )?; + Ok(true) + } else { + Ok(false) + } + } + + pub fn doc_search( + &self, + project_id: &str, + query: &str, + limit: usize, + ) -> rusqlite::Result> { + let mut stmt = self.conn.prepare( + "SELECT d.id, d.project_id, d.path, + snippet(store_docs_fts, 0, '', '', '...', 48) AS snip, + rank + FROM store_docs_fts + JOIN store_documents d ON d.rowid = store_docs_fts.rowid + WHERE store_docs_fts MATCH ?1 + AND d.project_id = ?2 + AND d.deleted_at IS NULL + ORDER BY rank + LIMIT ?3", + )?; + let rows = stmt.query_map(params![query, project_id, limit as i64], |row| { + Ok(DocMatch { + id: row.get(0)?, + project_id: row.get(1)?, + path: row.get(2)?, + snippet: row.get::<_, String>(3).unwrap_or_default(), + score: -row.get::<_, f64>(4)?, + }) + })?; + rows.collect() + } + + pub fn doc_set_embedding(&self, doc_id: &str, embedding: &[f32]) -> rusqlite::Result { + let now = db_kit::ids::now(); + let bytes: &[u8] = bytemuck::cast_slice(embedding); + let n = self.conn.execute( + "INSERT INTO store_docs_vec (doc_id, embedding, updated_at) VALUES (?1, ?2, ?3) + ON CONFLICT(doc_id) DO UPDATE SET embedding = ?2, updated_at = ?3", + params![doc_id, bytes, now], + )?; + Ok(n > 0) + } + + pub fn doc_get_embedding(&self, doc_id: &str) -> rusqlite::Result>> { + self.conn + .query_row( + "SELECT embedding FROM store_docs_vec WHERE doc_id = ?1", + params![doc_id], + |row| { + let blob: Vec = row.get(0)?; + let vec: Vec = blob + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + Ok(vec) + }, + ) + .optional() + } + + pub fn doc_vec_search( + &self, + project_id: &str, + query_vec: &[f32], + limit: usize, + ) -> rusqlite::Result> { + let mut stmt = self.conn.prepare( + "SELECT d.id, d.project_id, d.path, v.embedding + FROM store_docs_vec v + JOIN store_documents d ON d.id = v.doc_id + WHERE d.project_id = ?1 AND d.deleted_at IS NULL", + )?; + let mut results: Vec<(f64, DocMatch)> = stmt + .query_map(params![project_id], |row| { + let id: String = row.get(0)?; + let project_id: String = row.get(1)?; + let path: String = row.get(2)?; + let blob: Vec = row.get(3)?; + Ok((id, project_id, path, blob)) + })? + .filter_map(|r| r.ok()) + .filter_map(|(id, pid, path, blob)| { + if blob.len() % 4 != 0 { + return None; + } + let doc_vec: Vec = blob + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + let sim = crate::embed::cosine_similarity(query_vec, &doc_vec) as f64; + Some(( + sim, + DocMatch { + id, + project_id: pid, + path, + snippet: String::new(), + score: sim, + }, + )) + }) + .collect(); + results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + results.truncate(limit); + Ok(results.into_iter().map(|(_, m)| m).collect()) + } + + pub fn doc_hybrid_search( + &self, + project_id: &str, + fts_query: &str, + query_vec: &[f32], + limit: usize, + alpha: f64, + ) -> rusqlite::Result> { + let mut fts = self.doc_search(project_id, fts_query, limit * 2)?; + let mut vec = self.doc_vec_search(project_id, query_vec, limit * 2)?; + + let mut max_fts = fts.first().map(|m| m.score).unwrap_or(1.0); + let mut max_vec = vec.first().map(|m| m.score).unwrap_or(1.0); + if max_fts < 1e-12 { + max_fts = 1.0; + } + if max_vec < 1e-12 { + max_vec = 1.0; + } + + for m in &mut fts { + m.score = alpha * (m.score / max_fts); + } + for m in &mut vec { + m.score = (1.0 - alpha) * (m.score / max_vec); + } + + let mut combined: Vec = Vec::new(); + let mut seen = std::collections::HashSet::new(); + + for m in fts.into_iter().chain(vec) { + if seen.insert(m.id.clone()) { + combined.push(m); + } else if let Some(existing) = combined.iter_mut().find(|e| e.id == m.id) { + existing.score += m.score; + } + } + + combined.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + combined.truncate(limit); + Ok(combined) + } + + pub fn doc_list(&self, project_id: &str) -> rusqlite::Result> { + let mut stmt = self.conn.prepare( + "SELECT id, project_id, path, content, created_at, updated_at, deleted_at + FROM store_documents + WHERE project_id = ?1 AND deleted_at IS NULL + ORDER BY path", + )?; + let rows = stmt.query_map(params![project_id], |row| { + Ok(Document { + id: row.get(0)?, + project_id: row.get(1)?, + path: row.get(2)?, + content: row.get(3)?, + created_at: row.get(4)?, + updated_at: row.get(5)?, + deleted_at: row.get(6)?, + }) + })?; + rows.collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn store() -> Store { + Store::open_memory().unwrap() + } + + #[test] + fn create_and_read() { + let s = store(); + let doc = s.doc_upsert("proj-1", "/hello.md", "Hello world").unwrap(); + assert_eq!(doc.project_id, "proj-1"); + assert_eq!(doc.path, "/hello.md"); + assert!(doc.deleted_at.is_none()); + + let fetched = s.doc_get(&doc.id).unwrap().unwrap(); + assert_eq!(fetched.content, "Hello world"); + } + + #[test] + fn upsert_updates_existing() { + let s = store(); + let doc = s.doc_upsert("p", "/same.md", "v1").unwrap(); + let updated = s.doc_upsert("p", "/same.md", "v2").unwrap(); + assert_eq!(updated.id, doc.id); + assert_eq!(updated.content, "v2"); + } + + #[test] + fn soft_delete_and_list() { + let s = store(); + s.doc_upsert("p", "/a.md", "a").unwrap(); + let b = s.doc_upsert("p", "/b.md", "b").unwrap(); + assert_eq!(s.doc_list("p").unwrap().len(), 2); + + s.doc_delete(&b.id).unwrap(); + let list = s.doc_list("p").unwrap(); + assert_eq!(list.len(), 1); + assert_eq!(list[0].path, "/a.md"); + } + + #[test] + fn fts_search_finds_matching_content() { + let s = store(); + s.doc_upsert("p", "/rust.md", "Rust is a systems programming language") + .unwrap(); + s.doc_upsert("p", "/go.md", "Go is fast and concurrent") + .unwrap(); + s.doc_upsert("p", "/python.md", "Python is great for data science") + .unwrap(); + + let results = s.doc_search("p", "rust", 10).unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].path, "/rust.md"); + assert!(results[0].score > 0.0); + } + + #[test] + fn fts_search_multi_word() { + let s = store(); + s.doc_upsert("p", "/a.md", "the quick brown fox").unwrap(); + s.doc_upsert("p", "/b.md", "jumps over the lazy dog") + .unwrap(); + + let results = s.doc_search("p", "quick fox", 10).unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].path, "/a.md"); + } + + #[test] + fn fts_search_scoped_to_project() { + let s = store(); + s.doc_upsert("p1", "/doc.md", "shared term").unwrap(); + s.doc_upsert("p2", "/doc.md", "shared term").unwrap(); + + let results = s.doc_search("p1", "shared term", 10).unwrap(); + assert_eq!(results.len(), 1); + } + + fn test_embed(s: &Store, doc_id: &str, dim: usize, val: f32) { + let embedding = vec![val; dim]; + s.doc_set_embedding(doc_id, &embedding).unwrap(); + } + + #[test] + fn set_and_get_embedding() { + let s = store(); + let doc = s.doc_upsert("p", "/doc.md", "content").unwrap(); + let emb = vec![0.1, 0.2, 0.3]; + s.doc_set_embedding(&doc.id, &emb).unwrap(); + let got = s.doc_get_embedding(&doc.id).unwrap().unwrap(); + assert_eq!(got.len(), 3); + assert!((got[0] - 0.1).abs() < 1e-6); + } + + #[test] + fn vec_search_ranks_by_similarity() { + let s = store(); + let d1 = s.doc_upsert("p", "/cat.md", "about cats").unwrap(); + let d2 = s.doc_upsert("p", "/dog.md", "about dogs").unwrap(); + let d3 = s.doc_upsert("p", "/car.md", "about cars").unwrap(); + + test_embed(&s, &d1.id, 4, 1.0); + test_embed(&s, &d2.id, 4, 0.8); + test_embed(&s, &d3.id, 4, 0.0); + + let query = vec![1.0; 4]; + let results = s.doc_vec_search("p", &query, 10).unwrap(); + assert_eq!(results.len(), 3); + assert_eq!(results[0].id, d1.id); + assert_eq!(results[1].id, d2.id); + } + + #[test] + fn hybrid_search_combines_scores() { + let s = store(); + let d1 = s + .doc_upsert("p", "/rust.md", "Rust programming language") + .unwrap(); + s.doc_upsert("p", "/other.md", "Something else entirely") + .unwrap(); + + test_embed(&s, &d1.id, 4, 1.0); + + let query_vec = vec![1.0; 4]; + let results = s + .doc_hybrid_search("p", "rust", &query_vec, 10, 0.5) + .unwrap(); + assert!(!results.is_empty()); + assert_eq!(results[0].id, d1.id); + } +} diff --git a/crates/agentflare-store/src/embed.rs b/crates/agentflare-store/src/embed.rs new file mode 100644 index 00000000..8d5dcbea --- /dev/null +++ b/crates/agentflare-store/src/embed.rs @@ -0,0 +1,50 @@ +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); + let na: f32 = a.iter().map(|x| x * x).sum(); + let nb: f32 = b.iter().map(|x| x * x).sum(); + let denom = na.sqrt() * nb.sqrt(); + if denom < 1e-12 { 0.0 } else { dot / denom } +} + +pub fn normalize(v: &mut [f32]) { + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 1e-12 { + for x in v.iter_mut() { + *x /= norm; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cosine_similarity_identical() { + let a = vec![1.0, 0.0, 0.0]; + let b = vec![1.0, 0.0, 0.0]; + assert!((cosine_similarity(&a, &b) - 1.0).abs() < 1e-6); + } + + #[test] + fn cosine_similarity_orthogonal() { + let a = vec![1.0, 0.0]; + let b = vec![0.0, 1.0]; + assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6); + } + + #[test] + fn cosine_similarity_zero_denom() { + let a = vec![0.0, 0.0]; + let b = vec![1.0, 0.0]; + assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6); + } + + #[test] + fn normalize_unit_vector() { + let mut v = vec![3.0, 4.0]; + normalize(&mut v); + let len: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + assert!((len - 1.0).abs() < 1e-6); + } +} diff --git a/crates/agentflare-store/src/kv.rs b/crates/agentflare-store/src/kv.rs new file mode 100644 index 00000000..ea1c77e3 --- /dev/null +++ b/crates/agentflare-store/src/kv.rs @@ -0,0 +1,130 @@ +use crate::Store; +use rusqlite::OptionalExtension; +use rusqlite::params; +use serde::Serialize; + +#[derive(Debug, Clone, Serialize)] +pub struct KvEntry { + pub key: String, + pub value: Vec, + pub created_at: i64, + pub updated_at: i64, +} + +impl Store { + pub fn kv_set(&self, key: &str, value: &[u8]) -> rusqlite::Result<()> { + let now = db_kit::ids::now(); + self.conn.execute( + "INSERT INTO store_kv (key, value, created_at, updated_at) + VALUES (?1, ?2, ?3, ?3) + ON CONFLICT(key) DO UPDATE SET value = ?2, updated_at = ?3", + params![key, value, now], + )?; + Ok(()) + } + + pub fn kv_get(&self, key: &str) -> rusqlite::Result> { + self.conn + .query_row( + "SELECT key, value, created_at, updated_at FROM store_kv WHERE key = ?1", + params![key], + |row| { + Ok(KvEntry { + key: row.get(0)?, + value: row.get(1)?, + created_at: row.get(2)?, + updated_at: row.get(3)?, + }) + }, + ) + .optional() + } + + pub fn kv_exists(&self, key: &str) -> rusqlite::Result { + self.conn + .query_row( + "SELECT 1 FROM store_kv WHERE key = ?1", + params![key], + |_| Ok(()), + ) + .optional() + .map(|o| o.is_some()) + } + + pub fn kv_delete(&self, key: &str) -> rusqlite::Result { + let n = self + .conn + .execute("DELETE FROM store_kv WHERE key = ?1", params![key])?; + Ok(n > 0) + } + + pub fn kv_scan(&self, prefix: &str) -> rusqlite::Result> { + let mut stmt = self + .conn + .prepare("SELECT key, value, created_at, updated_at FROM store_kv WHERE key >= ?1 AND key < ?2 ORDER BY key")?; + let end = { + let mut s = prefix.to_string(); + s.push('\u{10FFFF}'); + s + }; + let rows = stmt.query_map(params![prefix, end], |row| { + Ok(KvEntry { + key: row.get(0)?, + value: row.get(1)?, + created_at: row.get(2)?, + updated_at: row.get(3)?, + }) + })?; + rows.collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn store() -> Store { + Store::open_memory().unwrap() + } + + #[test] + fn set_and_get() { + let s = store(); + s.kv_set("hello", b"world").unwrap(); + let entry = s.kv_get("hello").unwrap().unwrap(); + assert_eq!(entry.key, "hello"); + assert_eq!(entry.value, b"world"); + } + + #[test] + fn get_missing() { + let s = store(); + assert!(s.kv_get("nope").unwrap().is_none()); + } + + #[test] + fn exists() { + let s = store(); + assert!(!s.kv_exists("x").unwrap()); + s.kv_set("x", b"1").unwrap(); + assert!(s.kv_exists("x").unwrap()); + } + + #[test] + fn delete() { + let s = store(); + s.kv_set("x", b"1").unwrap(); + assert!(s.kv_delete("x").unwrap()); + assert!(!s.kv_exists("x").unwrap()); + } + + #[test] + fn scan_prefix() { + let s = store(); + s.kv_set("a:1", b"").unwrap(); + s.kv_set("a:2", b"").unwrap(); + s.kv_set("b:1", b"").unwrap(); + let entries = s.kv_scan("a:").unwrap(); + assert_eq!(entries.len(), 2); + } +} diff --git a/crates/agentflare-store/src/leases.rs b/crates/agentflare-store/src/leases.rs new file mode 100644 index 00000000..db127018 --- /dev/null +++ b/crates/agentflare-store/src/leases.rs @@ -0,0 +1,55 @@ +use db_kit::claim::ClaimLedger; +use rusqlite::Connection; + +pub static LEDGER: ClaimLedger = ClaimLedger::new("store_leases", &["key"]); + +pub fn acquire( + conn: &Connection, + key: &str, + owner: &str, + ttl_secs: i64, +) -> rusqlite::Result { + let now = db_kit::ids::now(); + LEDGER.acquire(conn, &[key], owner, now, ttl_secs) +} + +pub use db_kit::claim::Acquire; + +pub fn heartbeat(conn: &Connection, key: &str, owner: &str) -> rusqlite::Result { + let now = db_kit::ids::now(); + LEDGER.heartbeat(conn, &[key], owner, now) +} + +pub fn release(conn: &Connection, key: &str, owner: &str) -> rusqlite::Result { + LEDGER.release(conn, &[key], owner) +} + +pub fn is_owner(conn: &Connection, key: &str, owner: &str) -> rusqlite::Result { + LEDGER.is_owner(conn, &[key], owner) +} + +pub fn done(conn: &Connection, key: &str, owner: &str) -> rusqlite::Result { + let now = db_kit::ids::now(); + LEDGER.done(conn, &[key], owner, now) +} + +#[cfg(test)] +mod tests { + use super::*; + use db_kit::open_file; + + #[test] + fn acquire_and_release_cycle() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("test.db"); + let conn = open_file(&path, &super::super::migrations::migrations()).unwrap(); + + let key = "test-key"; + let owner = "test-owner"; + + assert_eq!(acquire(&conn, key, owner, 60).unwrap(), Acquire::Acquired); + assert!(heartbeat(&conn, key, owner).unwrap()); + assert!(release(&conn, key, owner).unwrap()); + assert!(!is_owner(&conn, key, owner).unwrap()); + } +} diff --git a/crates/agentflare-store/src/lib.rs b/crates/agentflare-store/src/lib.rs new file mode 100644 index 00000000..41d4d8e8 --- /dev/null +++ b/crates/agentflare-store/src/lib.rs @@ -0,0 +1,74 @@ +pub mod blobs; +pub mod documents; +pub mod embed; +pub mod kv; +pub mod leases; +pub mod migrate; +pub mod migrations; + +use rusqlite::Connection; +use std::path::{Path, PathBuf}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Sqlite(#[from] rusqlite::Error), + #[error(transparent)] + Migration(#[from] rusqlite_migration::Error), + #[error(transparent)] + DbKit(#[from] db_kit::open::Error), + #[error("not found: {0}")] + NotFound(String), + #[error("blob too large: {size} > {max}")] + BlobTooLarge { size: u64, max: u64 }, + #[error("lease denied: {0}")] + LeaseDenied(String), +} + +pub struct Store { + conn: Connection, + root: PathBuf, +} + +impl Store { + pub fn open_file(path: &Path) -> Result { + let conn = db_kit::open_file(path, &migrations::migrations())?; + let root = path.parent().unwrap_or(Path::new(".")).to_path_buf(); + Ok(Self { conn, root }) + } + + pub fn open_memory() -> Result { + let conn = db_kit::open_memory(&migrations::migrations())?; + Ok(Self { + conn, + root: PathBuf::from(":memory:"), + }) + } + + pub fn conn(&self) -> &Connection { + &self.conn + } + + pub fn root(&self) -> &Path { + &self.root + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn open_memory_store() { + let store = Store::open_memory().unwrap(); + store.conn().execute_batch("SELECT 1").unwrap(); + } + + #[test] + fn open_file_store() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("store.db"); + let store = Store::open_file(&path).unwrap(); + assert!(store.root().exists()); + } +} diff --git a/crates/agentflare-store/src/migrate.rs b/crates/agentflare-store/src/migrate.rs new file mode 100644 index 00000000..313d6df2 --- /dev/null +++ b/crates/agentflare-store/src/migrate.rs @@ -0,0 +1,95 @@ +use crate::Store; +use std::path::Path; + +const MIGRATION_MARKER: &str = "_migration_state_json_v1"; + +#[derive(Debug, thiserror::Error)] +pub enum MigrateError { + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + #[error("JSON error: {0}")] + Json(#[from] serde_json::Error), + #[error("Store error: {0}")] + Store(#[from] rusqlite::Error), + #[error("migration already ran at {0}")] + AlreadyRan(i64), +} + +pub fn migrate_state_json(store: &Store, path: &Path) -> Result { + if let Some(existing) = store.kv_get(MIGRATION_MARKER)? { + let ts: i64 = serde_json::from_slice(&existing.value)?; + return Err(MigrateError::AlreadyRan(ts)); + } + + let content = std::fs::read_to_string(path)?; + let map: serde_json::Map = serde_json::from_str(&content)?; + + for (key, value) in &map { + let serialized = serde_json::to_vec(value)?; + store.kv_set(key, &serialized)?; + } + + let now = db_kit::ids::now(); + let marker = serde_json::to_vec(&now)?; + store.kv_set(MIGRATION_MARKER, &marker)?; + + Ok(map.len()) +} + +pub fn validate_roundtrip(store: &Store, path: &Path) -> Result { + let content = std::fs::read_to_string(path)?; + let original: serde_json::Map = serde_json::from_str(&content)?; + + for (key, expected) in &original { + let entry = store.kv_get(key)?; + match entry { + None => return Ok(false), + Some(entry) => { + let stored: serde_json::Value = serde_json::from_slice(&entry.value)?; + if &stored != expected { + return Ok(false); + } + } + } + } + + Ok(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Store; + + #[test] + fn migrate_and_validate() { + let dir = tempfile::tempdir().unwrap(); + let state_path = dir.path().join("state.json"); + std::fs::write(&state_path, r#"{"hello":"world","count":42,"flag":true}"#).unwrap(); + + let db_path = dir.path().join("store.db"); + let store = Store::open_file(&db_path).unwrap(); + + let count = migrate_state_json(&store, &state_path).unwrap(); + assert_eq!(count, 3); + + let hello = store.kv_get("hello").unwrap().unwrap(); + let val: serde_json::Value = serde_json::from_slice(&hello.value).unwrap(); + assert_eq!(val, serde_json::json!("world")); + + assert!(validate_roundtrip(&store, &state_path).unwrap()); + } + + #[test] + fn rejects_repeat_migration() { + let dir = tempfile::tempdir().unwrap(); + let state_path = dir.path().join("state.json"); + std::fs::write(&state_path, r#"{"k":"v"}"#).unwrap(); + + let store = Store::open_memory().unwrap(); + migrate_state_json(&store, &state_path).unwrap(); + + let err = migrate_state_json(&store, &state_path).unwrap_err(); + assert!(matches!(err, MigrateError::AlreadyRan(_))); + } +} diff --git a/crates/agentflare-store/src/migrations.rs b/crates/agentflare-store/src/migrations.rs new file mode 100644 index 00000000..67b3e17a --- /dev/null +++ b/crates/agentflare-store/src/migrations.rs @@ -0,0 +1,61 @@ +use rusqlite_migration::{M, Migrations}; + +pub fn migrations() -> Migrations<'static> { + Migrations::new(vec![ + M::up( + "CREATE TABLE IF NOT EXISTS store_kv ( + key TEXT PRIMARY KEY NOT NULL, + value BLOB NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + );", + ), + M::up( + "CREATE TABLE IF NOT EXISTS store_documents ( + id TEXT PRIMARY KEY NOT NULL, + project_id TEXT NOT NULL DEFAULT '', + path TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + deleted_at INTEGER + ); + CREATE INDEX IF NOT EXISTS idx_docs_project ON store_documents(project_id);", + ), + M::up( + "CREATE VIRTUAL TABLE IF NOT EXISTS store_docs_fts USING fts5( + content + );", + ), + M::up( + "CREATE TABLE IF NOT EXISTS store_docs_vec ( + doc_id TEXT PRIMARY KEY NOT NULL, + embedding BLOB NOT NULL, + updated_at INTEGER NOT NULL + );", + ), + M::up( + "CREATE TABLE IF NOT EXISTS store_blobs ( + hash TEXT PRIMARY KEY NOT NULL, + size INTEGER NOT NULL, + ref_count INTEGER NOT NULL DEFAULT 1, + created_at INTEGER NOT NULL + ); + CREATE TABLE IF NOT EXISTS store_blob_chunks ( + hash TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + data BLOB NOT NULL, + PRIMARY KEY (hash, chunk_index) + );", + ), + M::up( + "CREATE TABLE IF NOT EXISTS store_leases ( + key TEXT PRIMARY KEY NOT NULL, + owner TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'claimed', + created_at INTEGER NOT NULL, + heartbeat_at INTEGER NOT NULL + );", + ), + ]) +} From 7eef3635176f20349927609ffbf0b0aef4559d36 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 14:40:05 +0530 Subject: [PATCH 04/16] feat(agentflare-store): initial crate - kv, documents (CRUD+FTS+vector), blobs, leases --- Cargo.lock | 358 +++++++++++++++++- Cargo.toml | 1 + crates/agentflare-store/Cargo.toml | 10 + crates/agentflare-store/src/blobs.rs | 122 ++++-- crates/agentflare-store/src/documents.rs | 357 +++++++++++++---- .../src/embedding_pipeline/download.rs | 177 +++++++++ .../src/embedding_pipeline/mod.rs | 297 +++++++++++++++ .../src/embedding_pipeline/model_registry.rs | 200 ++++++++++ .../src/embedding_pipeline/pooling.rs | 41 ++ .../src/embedding_pipeline/tokenizer.rs | 204 ++++++++++ crates/agentflare-store/src/kv.rs | 21 +- crates/agentflare-store/src/lib.rs | 16 +- crates/agentflare-store/src/migrations.rs | 23 ++ src/main.rs | 1 + src/store.rs | 36 ++ 15 files changed, 1740 insertions(+), 124 deletions(-) create mode 100644 crates/agentflare-store/src/embedding_pipeline/download.rs create mode 100644 crates/agentflare-store/src/embedding_pipeline/mod.rs create mode 100644 crates/agentflare-store/src/embedding_pipeline/model_registry.rs create mode 100644 crates/agentflare-store/src/embedding_pipeline/pooling.rs create mode 100644 crates/agentflare-store/src/embedding_pipeline/tokenizer.rs create mode 100644 src/store.rs diff --git a/Cargo.lock b/Cargo.lock index e6311a19..caaa5581 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -75,6 +75,7 @@ dependencies = [ "agentflare-flare-output", "agentflare-gateway-registry", "agentflare-skill-registry", + "agentflare-store", "base64", "built", "chrono", @@ -103,7 +104,7 @@ dependencies = [ "tempfile", "thiserror", "tokio", - "ureq", + "ureq 2.12.1", "zip", ] @@ -142,7 +143,7 @@ dependencies = [ "sha2", "tempfile", "thiserror", - "ureq", + "ureq 2.12.1", "url", "uuid", ] @@ -178,7 +179,7 @@ dependencies = [ "serde_json", "tempfile", "thiserror", - "ureq", + "ureq 2.12.1", "which", ] @@ -202,7 +203,7 @@ dependencies = [ "tokio", "tokio-util", "toml", - "ureq", + "ureq 2.12.1", ] [[package]] @@ -226,15 +227,23 @@ name = "agentflare-store" version = "0.1.0" dependencies = [ "agentflare-db-kit", + "anyhow", "blake3", "bytemuck", + "dirs", + "ndarray", + "ort", "parking_lot", + "rayon", "rusqlite", "rusqlite_migration", "serde", "serde_json", + "sha2", "tempfile", "thiserror", + "tracing", + "ureq 2.12.1", ] [[package]] @@ -305,6 +314,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + [[package]] name = "arbitrary" version = "1.4.2" @@ -486,6 +501,12 @@ dependencies = [ "syn", ] +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.12.0" @@ -667,6 +688,16 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -700,6 +731,25 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" @@ -760,6 +810,16 @@ dependencies = [ "syn", ] +[[package]] +name = "der" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d" +dependencies = [ + "pem-rfc7468", + "zeroize", +] + [[package]] name = "derive_arbitrary" version = "1.4.2" @@ -921,6 +981,21 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1115,6 +1190,12 @@ dependencies = [ "digest", ] +[[package]] +name = "hmac-sha256" +version = "1.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" + [[package]] name = "home" version = "0.5.12" @@ -1491,12 +1572,28 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "lzma-rust2" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e20f57f9918e5bd7bc58c22cdd70a6afc7375d4dd9683af5f2b34bd3d2bba619" + [[package]] name = "matchit" version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.2" @@ -1530,6 +1627,38 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "native-tls" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "nix" version = "0.31.3" @@ -1551,6 +1680,24 @@ dependencies = [ "winapi", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1606,12 +1753,79 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" +[[package]] +name = "openssl" +version = "0.10.81" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77823a27f0babb03091cb9ed9ef80af3b39dbc82f97e8fa530374b7dafd87a45" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "openssl-sys" +version = "0.9.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b47e7e6bb2c38cd930d25a23b40fa52e068c10e85f3e03a7f5ba5aaca5713695" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "option-ext" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "ort" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec", + "tracing", + "ureq 3.3.0", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" +dependencies = [ + "hmac-sha256", + "lzma-rust2", + "ureq 3.3.0", +] + [[package]] name = "owo-colors" version = "4.3.0" @@ -1670,6 +1884,15 @@ dependencies = [ "sha2", ] +[[package]] +name = "pem-rfc7468" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1702,9 +1925,18 @@ dependencies = [ [[package]] name = "portable-atomic" -version = "1.14.0" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] [[package]] name = "potential_utf" @@ -1819,6 +2051,32 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2117,6 +2375,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "schemars" version = "1.2.1" @@ -2149,6 +2416,29 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "serde" version = "1.0.228" @@ -2301,6 +2591,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + [[package]] name = "sqlite-wasm-rs" version = "0.5.5" @@ -2742,6 +3043,36 @@ dependencies = [ "webpki-roots 0.26.11", ] +[[package]] +name = "ureq" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" +dependencies = [ + "base64", + "der", + "log", + "native-tls", + "percent-encoding", + "rustls-pki-types", + "socks", + "ureq-proto", + "utf8-zero", + "webpki-root-certs", +] + +[[package]] +name = "ureq-proto" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" +dependencies = [ + "base64", + "http", + "httparse", + "log", +] + [[package]] name = "url" version = "2.5.8" @@ -2754,6 +3085,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf8-zero" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -2898,6 +3235,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "0.26.11" diff --git a/Cargo.toml b/Cargo.toml index 0ee28420..71be4dda 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -72,6 +72,7 @@ agentflare-backend = { package = "agentflare-backend", path = "crates/agentflare db_kit = { package = "agentflare-db-kit", path = "crates/agentflare-db-kit" } agent-detector = "0.2.1" flare-search-kit = { path = "crates/flare-search-kit" } +agentflare-store = { path = "crates/agentflare-store" } [target.'cfg(unix)'.dependencies] libc = "0.2" diff --git a/crates/agentflare-store/Cargo.toml b/crates/agentflare-store/Cargo.toml index f4ec889e..1aa58f72 100644 --- a/crates/agentflare-store/Cargo.toml +++ b/crates/agentflare-store/Cargo.toml @@ -17,8 +17,18 @@ blake3 = "1.5" thiserror = "2" parking_lot = "0.12" bytemuck = { version = "1", features = ["derive"] } +anyhow = { version = "1", optional = true } +dirs = { version = "6", optional = true } +tracing = { version = "0.1", optional = true } +ureq = { version = "2", features = ["json"], optional = true } +sha2 = { version = "0.10", optional = true } +ndarray = { version = "0.17", optional = true } +rayon = { version = "1", optional = true } +ort = { version = "=2.0.0-rc.12", optional = true, default-features = true, features = ["ndarray"] } + [features] default = [] +embeddings = ["dep:anyhow", "dep:dirs", "dep:tracing", "dep:ureq", "dep:sha2", "dep:ndarray", "dep:rayon", "dep:ort"] [dev-dependencies] tempfile = "3" diff --git a/crates/agentflare-store/src/blobs.rs b/crates/agentflare-store/src/blobs.rs index 889fc156..fe1c81c1 100644 --- a/crates/agentflare-store/src/blobs.rs +++ b/crates/agentflare-store/src/blobs.rs @@ -1,6 +1,7 @@ use crate::Store; use rusqlite::{OptionalExtension, params}; use serde::Serialize; +use std::path::Path; #[derive(Debug, Clone, Serialize)] pub struct BlobMeta { @@ -10,16 +11,46 @@ pub struct BlobMeta { pub created_at: i64, } -const CHUNK_SIZE: usize = 64 * 1024; // 64 KiB +const CHUNK_SIZE: usize = 64 * 1024; + +fn blob_disk_path(root: &Path, hash: &str) -> PathBuf { + root.join("blobs").join(&hash[..2]).join(hash) +} + +fn read_disk_blob(root: &Path, hash: &str) -> Option> { + let path = blob_disk_path(root, hash); + std::fs::read(&path).ok() +} + +fn write_disk_blob(root: &Path, hash: &str, data: &[u8]) -> Result<(), std::io::Error> { + let path = blob_disk_path(root, hash); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(&path, data)?; + std::fs::write(path.with_extension("meta"), &[])?; + Ok(()) +} + +fn delete_disk_blob(root: &Path, hash: &str) { + let path = blob_disk_path(root, hash); + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(path.with_extension("meta")); +} + +use std::path::PathBuf; impl Store { + fn is_memory(&self) -> bool { + self.root.to_string_lossy() == ":memory:" + } + pub fn blob_store(&self, data: &[u8]) -> rusqlite::Result { + let conn = self.conn(); let hash = blake3::hash(data).to_hex().to_string(); let now = db_kit::ids::now(); - // Bump ref count if exists - let exists = self - .conn + let exists = conn .query_row( "SELECT 1 FROM store_blobs WHERE hash = ?1", params![hash], @@ -29,31 +60,38 @@ impl Store { .is_some(); if exists { - self.conn.execute( + conn.execute( "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", params![hash], )?; return Ok(hash); } - self.conn.execute( + let is_memory = self.is_memory(); + if !is_memory { + if let Err(e) = write_disk_blob(&self.root, &hash, data) { + return Err(rusqlite::Error::ToSqlConversionFailure(Box::new(e))); + } + } else { + for (i, chunk) in data.chunks(CHUNK_SIZE).enumerate() { + conn.execute( + "INSERT INTO store_blob_chunks (hash, chunk_index, data) VALUES (?1, ?2, ?3)", + params![hash, i as i64, chunk], + )?; + } + } + + conn.execute( "INSERT INTO store_blobs (hash, size, ref_count, created_at) VALUES (?1, ?2, 1, ?3)", params![hash, data.len() as i64, now], )?; - - for (i, chunk) in data.chunks(CHUNK_SIZE).enumerate() { - self.conn.execute( - "INSERT INTO store_blob_chunks (hash, chunk_index, data) VALUES (?1, ?2, ?3)", - params![hash, i as i64, chunk], - )?; - } - Ok(hash) } pub fn blob_get(&self, hash: &str) -> rusqlite::Result>> { let meta: BlobMeta = match self .conn + .lock() .query_row( "SELECT hash, size, ref_count, created_at FROM store_blobs WHERE hash = ?1", params![hash], @@ -72,8 +110,12 @@ impl Store { None => return Ok(None), }; - let mut stmt = self - .conn + if !self.is_memory() { + return Ok(read_disk_blob(&self.root, hash)); + } + + let conn = self.conn(); + let mut stmt = conn .prepare("SELECT data FROM store_blob_chunks WHERE hash = ?1 ORDER BY chunk_index")?; let chunks: rusqlite::Result>> = stmt.query_map(params![hash], |row| row.get(0))?.collect(); @@ -86,7 +128,8 @@ impl Store { } pub fn blob_ref(&self, hash: &str) -> rusqlite::Result { - let n = self.conn.execute( + let conn = self.conn(); + let n = conn.execute( "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", params![hash], )?; @@ -94,19 +137,31 @@ impl Store { } pub fn blob_unref(&self, hash: &str) -> rusqlite::Result { - let n = self.conn.execute( + let conn = self.conn(); + let n = conn.execute( "UPDATE store_blobs SET ref_count = ref_count - 1 WHERE hash = ?1 AND ref_count > 0", params![hash], )?; if n > 0 { - self.conn.execute( - "DELETE FROM store_blobs WHERE hash = ?1 AND ref_count <= 0", - params![hash], - )?; - self.conn.execute( - "DELETE FROM store_blob_chunks WHERE hash = ?1", - params![hash], - )?; + let removed = conn + .query_row( + "SELECT ref_count <= 0 FROM store_blobs WHERE hash = ?1", + params![hash], + |row| row.get::<_, bool>(0), + ) + .optional()? + .unwrap_or(false); + if removed { + conn.execute("DELETE FROM store_blobs WHERE hash = ?1", params![hash])?; + if !self.is_memory() { + delete_disk_blob(&self.root, hash); + } else { + conn.execute( + "DELETE FROM store_blob_chunks WHERE hash = ?1", + params![hash], + )?; + } + } } Ok(n > 0) } @@ -154,4 +209,19 @@ mod tests { assert!(s.blob_unref(&h).unwrap()); assert!(s.blob_get(&h).unwrap().is_none()); } + + #[test] + fn disk_storage() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("store.db"); + let s = Store::open_file(&db_path).unwrap(); + let data = b"content-addressed on disk"; + let hash = s.blob_store(data).unwrap(); + + let disk_path = blob_disk_path(dir.path(), &hash); + assert!(disk_path.exists(), "blob file should exist on disk"); + + let retrieved = s.blob_get(&hash).unwrap().unwrap(); + assert_eq!(retrieved, data); + } } diff --git a/crates/agentflare-store/src/documents.rs b/crates/agentflare-store/src/documents.rs index e5ef871f..9e141732 100644 --- a/crates/agentflare-store/src/documents.rs +++ b/crates/agentflare-store/src/documents.rs @@ -1,18 +1,38 @@ use crate::Store; use rusqlite::{OptionalExtension, params}; -use serde::Serialize; +use serde::{Deserialize, Serialize}; -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct Document { pub id: String, pub project_id: String, pub path: String, pub content: String, + pub title: String, + pub doc_type: String, + pub blob_hash: Option, + pub mime: String, + pub tags: Vec, + pub session_id: Option, + pub source: String, + pub version: i32, pub created_at: i64, pub updated_at: i64, pub deleted_at: Option, } +#[derive(Debug, Clone, Serialize)] +pub struct DocVersion { + pub id: String, + pub doc_id: String, + pub version: i32, + pub content: String, + pub blob_hash: Option, + pub mime: String, + pub title: String, + pub created_at: i64, +} + #[derive(Debug, Clone, Serialize)] pub struct DocMatch { pub id: String, @@ -22,58 +42,187 @@ pub struct DocMatch { pub score: f64, } +#[derive(Debug, Default)] +pub struct DocUpsertOpts { + pub title: Option, + pub doc_type: Option, + pub blob_hash: Option, + pub mime: Option, + pub tags: Option>, + pub session_id: Option, + pub source: Option, +} + impl Store { - fn doc_sync_fts(&self, row_id: i64, content: &str) -> rusqlite::Result<()> { - // FTS5 has no REPLACE/UPSERT — delete any existing rowid first (no-op if fresh) - self.conn.execute( + fn doc_sync_fts(conn: &rusqlite::Connection, row_id: i64, content: &str) -> rusqlite::Result<()> { + conn.execute( "DELETE FROM store_docs_fts WHERE rowid = ?1", params![row_id], )?; - self.conn.execute( + conn.execute( "INSERT INTO store_docs_fts(rowid, content) VALUES (?1, ?2)", params![row_id, content], )?; Ok(()) } + + fn row_to_document(row: &rusqlite::Row) -> rusqlite::Result { + let tags_str: String = row.get(8)?; + let tags: Vec = serde_json::from_str(&tags_str).unwrap_or_default(); + Ok(Document { + id: row.get(0)?, + project_id: row.get(1)?, + path: row.get(2)?, + content: row.get(3)?, + title: row.get(4)?, + doc_type: row.get(5)?, + blob_hash: row.get(6)?, + mime: row.get(7)?, + tags, + session_id: row.get(9)?, + source: row.get(10)?, + version: row.get(11)?, + created_at: row.get(12)?, + updated_at: row.get(13)?, + deleted_at: row.get(14)?, + }) + } + pub fn doc_upsert( &self, project_id: &str, path: &str, content: &str, ) -> rusqlite::Result { + self.doc_upsert_with_opts(project_id, path, content, DocUpsertOpts::default()) + } + + pub fn doc_upsert_with_opts( + &self, + project_id: &str, + path: &str, + content: &str, + opts: DocUpsertOpts, + ) -> rusqlite::Result { + let conn = self.conn(); let now = db_kit::ids::now(); - let id = db_kit::ids::new_id(); - // Try to find existing by (project_id, path), else insert fresh - let existing = self - .conn + let existing = conn .query_row( - "SELECT id, rowid FROM store_documents WHERE project_id = ?1 AND path = ?2", + "SELECT id, rowid, content, version FROM store_documents + WHERE project_id = ?1 AND path = ?2", params![project_id, path], - |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, i64>(1)?, + row.get::<_, String>(2)?, + row.get::<_, i32>(3)?, + )) + }, ) .optional()?; - if let Some((existing_id, rowid)) = existing { - self.conn.execute( - "UPDATE store_documents SET content = ?1, updated_at = ?2, deleted_at = NULL WHERE id = ?3", - params![content, now, existing_id], + if let Some((existing_id, rowid, old_content, old_version)) = existing { + let new_version = old_version + 1; + let history_id = db_kit::ids::new_id(); + + // Snapshot current version to history + conn.execute( + "INSERT INTO store_doc_history (id, doc_id, version, content, title, created_at) + VALUES (?1, ?2, ?3, ?4, (SELECT title FROM store_documents WHERE id = ?2), ?5)", + params![history_id, existing_id, old_version, old_content, now], )?; - self.doc_sync_fts(rowid, content)?; + + conn.execute( + "UPDATE store_documents SET + content = ?1, updated_at = ?2, deleted_at = NULL, + version = ?3 + WHERE id = ?4", + params![content, now, new_version, existing_id], + )?; + + // Apply optional updates (need separate UPDATE to avoid long SQL) + if let Some(title) = &opts.title { + conn.execute( + "UPDATE store_documents SET title = ?1 WHERE id = ?2", + params![title, existing_id], + )?; + } + if let Some(doc_type) = &opts.doc_type { + conn.execute( + "UPDATE store_documents SET doc_type = ?1 WHERE id = ?2", + params![doc_type, existing_id], + )?; + } + if opts.blob_hash.is_some() { + conn.execute( + "UPDATE store_documents SET blob_hash = ?1 WHERE id = ?2", + params![opts.blob_hash, existing_id], + )?; + } + if let Some(mime) = &opts.mime { + conn.execute( + "UPDATE store_documents SET mime = ?1 WHERE id = ?2", + params![mime, existing_id], + )?; + } + if let Some(tags) = &opts.tags { + let json = serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string()); + conn.execute( + "UPDATE store_documents SET tags = ?1 WHERE id = ?2", + params![json, existing_id], + )?; + } + if opts.session_id.is_some() { + conn.execute( + "UPDATE store_documents SET session_id = ?1 WHERE id = ?2", + params![opts.session_id, existing_id], + )?; + } + if let Some(source) = &opts.source { + conn.execute( + "UPDATE store_documents SET source = ?1 WHERE id = ?2", + params![source, existing_id], + )?; + } + + Self::doc_sync_fts(&conn, rowid, content)?; + drop(conn); self.doc_get(&existing_id).map(|o| o.unwrap()) } else { - self.conn.execute( - "INSERT INTO store_documents (id, project_id, path, content, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?5)", - params![id, project_id, path, content, now], + let id = db_kit::ids::new_id(); + let title = opts.title.unwrap_or_default(); + let doc_type = opts.doc_type.unwrap_or_else(|| "file".to_string()); + let mime = opts.mime.unwrap_or_default(); + let tags_val = opts.tags.unwrap_or_default(); + let tags_json = serde_json::to_string(&tags_val).unwrap_or_else(|_| "[]".to_string()); + let source = opts.source.unwrap_or_default(); + + conn.execute( + "INSERT INTO store_documents + (id, project_id, path, content, title, doc_type, blob_hash, mime, tags, session_id, source, version, created_at, updated_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, 1, ?12, ?12)", + params![ + id, project_id, path, content, title, doc_type, opts.blob_hash, + mime, tags_json, opts.session_id, source, now + ], )?; - let rowid = self.conn.last_insert_rowid(); - self.doc_sync_fts(rowid, content)?; + let rowid = conn.last_insert_rowid(); + Self::doc_sync_fts(&conn, rowid, content)?; Ok(Document { id, project_id: project_id.to_string(), path: path.to_string(), content: content.to_string(), + title, + doc_type, + blob_hash: opts.blob_hash, + mime, + tags: tags_val, + session_id: opts.session_id, + source, + version: 1, created_at: now, updated_at: now, deleted_at: None, @@ -82,30 +231,21 @@ impl Store { } pub fn doc_get(&self, id: &str) -> rusqlite::Result> { - self.conn - .query_row( - "SELECT id, project_id, path, content, created_at, updated_at, deleted_at + let conn = self.conn(); + conn.query_row( + "SELECT id, project_id, path, content, title, doc_type, blob_hash, mime, tags, + session_id, source, version, created_at, updated_at, deleted_at FROM store_documents WHERE id = ?1", params![id], - |row| { - Ok(Document { - id: row.get(0)?, - project_id: row.get(1)?, - path: row.get(2)?, - content: row.get(3)?, - created_at: row.get(4)?, - updated_at: row.get(5)?, - deleted_at: row.get(6)?, - }) - }, + Self::row_to_document, ) .optional() } pub fn doc_delete(&self, id: &str) -> rusqlite::Result { + let conn = self.conn(); let now = db_kit::ids::now(); - if let Some(rowid) = self - .conn + if let Some(rowid) = conn .query_row( "SELECT rowid FROM store_documents WHERE id = ?1", params![id], @@ -113,11 +253,11 @@ impl Store { ) .optional()? { - self.conn.execute( + conn.execute( "UPDATE store_documents SET deleted_at = ?1 WHERE id = ?2", params![now, id], )?; - self.conn.execute( + conn.execute( "DELETE FROM store_docs_fts WHERE rowid = ?1", params![rowid], )?; @@ -128,8 +268,8 @@ impl Store { } pub fn doc_hard_delete(&self, id: &str) -> rusqlite::Result { - if let Some(rowid) = self - .conn + let conn = self.conn(); + if let Some(rowid) = conn .query_row( "SELECT rowid FROM store_documents WHERE id = ?1", params![id], @@ -137,9 +277,8 @@ impl Store { ) .optional()? { - self.conn - .execute("DELETE FROM store_documents WHERE id = ?1", params![id])?; - self.conn.execute( + conn.execute("DELETE FROM store_documents WHERE id = ?1", params![id])?; + conn.execute( "DELETE FROM store_docs_fts WHERE rowid = ?1", params![rowid], )?; @@ -149,13 +288,59 @@ impl Store { } } + pub fn doc_history(&self, doc_id: &str) -> rusqlite::Result> { + let conn = self.conn(); + let mut stmt = conn.prepare( + "SELECT id, doc_id, version, content, blob_hash, mime, title, created_at + FROM store_doc_history + WHERE doc_id = ?1 + ORDER BY version DESC", + )?; + let rows = stmt.query_map(params![doc_id], |row| { + Ok(DocVersion { + id: row.get(0)?, + doc_id: row.get(1)?, + version: row.get(2)?, + content: row.get(3)?, + blob_hash: row.get(4)?, + mime: row.get(5)?, + title: row.get(6)?, + created_at: row.get(7)?, + }) + })?; + rows.collect() + } + + pub fn doc_get_version(&self, doc_id: &str, version: i32) -> rusqlite::Result> { + let conn = self.conn(); + conn.query_row( + "SELECT id, doc_id, version, content, blob_hash, mime, title, created_at + FROM store_doc_history WHERE doc_id = ?1 AND version = ?2", + params![doc_id, version], + |row| { + Ok(DocVersion { + id: row.get(0)?, + doc_id: row.get(1)?, + version: row.get(2)?, + content: row.get(3)?, + blob_hash: row.get(4)?, + mime: row.get(5)?, + title: row.get(6)?, + created_at: row.get(7)?, + }) + }, + ) + .optional() + } + pub fn doc_search( &self, project_id: &str, query: &str, limit: usize, ) -> rusqlite::Result> { - let mut stmt = self.conn.prepare( + let conn = self.conn(); + let mut stmt = conn.prepare( "SELECT d.id, d.project_id, d.path, snippet(store_docs_fts, 0, '', '', '...', 48) AS snip, rank @@ -180,9 +365,10 @@ impl Store { } pub fn doc_set_embedding(&self, doc_id: &str, embedding: &[f32]) -> rusqlite::Result { + let conn = self.conn(); let now = db_kit::ids::now(); let bytes: &[u8] = bytemuck::cast_slice(embedding); - let n = self.conn.execute( + let n = conn.execute( "INSERT INTO store_docs_vec (doc_id, embedding, updated_at) VALUES (?1, ?2, ?3) ON CONFLICT(doc_id) DO UPDATE SET embedding = ?2, updated_at = ?3", params![doc_id, bytes, now], @@ -191,8 +377,8 @@ impl Store { } pub fn doc_get_embedding(&self, doc_id: &str) -> rusqlite::Result>> { - self.conn - .query_row( + let conn = self.conn(); + conn.query_row( "SELECT embedding FROM store_docs_vec WHERE doc_id = ?1", params![doc_id], |row| { @@ -213,7 +399,8 @@ impl Store { query_vec: &[f32], limit: usize, ) -> rusqlite::Result> { - let mut stmt = self.conn.prepare( + let conn = self.conn(); + let mut stmt = conn.prepare( "SELECT d.id, d.project_id, d.path, v.embedding FROM store_docs_vec v JOIN store_documents d ON d.id = v.doc_id @@ -267,19 +454,11 @@ impl Store { let mut max_fts = fts.first().map(|m| m.score).unwrap_or(1.0); let mut max_vec = vec.first().map(|m| m.score).unwrap_or(1.0); - if max_fts < 1e-12 { - max_fts = 1.0; - } - if max_vec < 1e-12 { - max_vec = 1.0; - } + if max_fts < 1e-12 { max_fts = 1.0; } + if max_vec < 1e-12 { max_vec = 1.0; } - for m in &mut fts { - m.score = alpha * (m.score / max_fts); - } - for m in &mut vec { - m.score = (1.0 - alpha) * (m.score / max_vec); - } + for m in &mut fts { m.score = alpha * (m.score / max_fts); } + for m in &mut vec { m.score = (1.0 - alpha) * (m.score / max_vec); } let mut combined: Vec = Vec::new(); let mut seen = std::collections::HashSet::new(); @@ -302,23 +481,15 @@ impl Store { } pub fn doc_list(&self, project_id: &str) -> rusqlite::Result> { - let mut stmt = self.conn.prepare( - "SELECT id, project_id, path, content, created_at, updated_at, deleted_at + let conn = self.conn(); + let mut stmt = conn.prepare( + "SELECT id, project_id, path, content, title, doc_type, blob_hash, mime, tags, + session_id, source, version, created_at, updated_at, deleted_at FROM store_documents WHERE project_id = ?1 AND deleted_at IS NULL ORDER BY path", )?; - let rows = stmt.query_map(params![project_id], |row| { - Ok(Document { - id: row.get(0)?, - project_id: row.get(1)?, - path: row.get(2)?, - content: row.get(3)?, - created_at: row.get(4)?, - updated_at: row.get(5)?, - deleted_at: row.get(6)?, - }) - })?; + let rows = stmt.query_map(params![project_id], Self::row_to_document)?; rows.collect() } } @@ -455,4 +626,38 @@ mod tests { assert!(!results.is_empty()); assert_eq!(results[0].id, d1.id); } + + #[test] + fn upsert_with_metadata() { + let s = store(); + let doc = s.doc_upsert_with_opts("p", "/meta.md", "content", DocUpsertOpts { + title: Some("My Doc".into()), + doc_type: Some("note".into()), + mime: Some("text/markdown".into()), + tags: Some(vec!["rust".into(), "db".into()]), + source: Some("agent".into()), + ..Default::default() + }).unwrap(); + assert_eq!(doc.title, "My Doc"); + assert_eq!(doc.doc_type, "note"); + assert_eq!(doc.mime, "text/markdown"); + assert_eq!(doc.tags, vec!["rust", "db"]); + assert_eq!(doc.source, "agent"); + assert_eq!(doc.version, 1); + } + + #[test] + fn versioning_increments_on_upsert() { + let s = store(); + let doc = s.doc_upsert("p", "/v.md", "v1").unwrap(); + assert_eq!(doc.version, 1); + + let updated = s.doc_upsert("p", "/v.md", "v2").unwrap(); + assert_eq!(updated.version, 2); + + let history = s.doc_history(&updated.id).unwrap(); + assert_eq!(history.len(), 1); + assert_eq!(history[0].version, 1); + assert_eq!(history[0].content, "v1"); + } } diff --git a/crates/agentflare-store/src/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs new file mode 100644 index 00000000..bd4859fd --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -0,0 +1,177 @@ +use std::collections::BTreeMap; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use super::model_registry::{ModelConfig, VocabSource}; + +const USER_AGENT: &str = concat!("agentflare-store/", env!("CARGO_PKG_VERSION")); +const LOCKFILE: &str = "model.lock.json"; + +struct DownloadFile { + url: String, + local_name: String, + min_bytes: u64, +} + +pub fn ensure_model(model_dir: &Path, config: &ModelConfig) -> anyhow::Result { + let files = download_files(config); + let all_present = files.iter().all(|f| model_dir.join(&f.local_name).exists()); + + if all_present { + return Ok(model_dir.to_path_buf()); + } + + tracing::info!( + "Embedding model '{}' not found, downloading to {}", + config.name, + model_dir.display() + ); + std::fs::create_dir_all(model_dir)?; + + let mut lock = read_lockfile(model_dir); + + for file in &files { + let local_path = model_dir.join(&file.local_name); + if local_path.exists() { + let meta = std::fs::metadata(&local_path)?; + if meta.len() >= file.min_bytes { + continue; + } + } + download_file(&file.url, &file.local_name, file.min_bytes, model_dir)?; + let actual = sha256_file(&model_dir.join(&file.local_name))?; + match lock.get(&file.local_name) { + Some(pinned) if pinned != &actual => { + let _ = std::fs::remove_file(model_dir.join(&file.local_name)); + anyhow::bail!( + "SHA-256 mismatch for {} of model '{}': pinned {pinned}, got {actual}. \ + Upstream content changed under same revision. Delete {} and re-download.", + file.local_name, config.name, + model_dir.join(LOCKFILE).display() + ); + } + Some(_) => {} + None => { + lock.insert(file.local_name.clone(), actual); + } + } + } + + write_lockfile(model_dir, &lock)?; + tracing::info!("Embedding model '{}' ready at {}", config.name, model_dir.display()); + Ok(model_dir.to_path_buf()) +} + +fn download_files(config: &ModelConfig) -> Vec { + vec![ + DownloadFile { + url: config.model_url(), + local_name: "model.onnx".to_string(), + min_bytes: config.model_min_bytes, + }, + DownloadFile { + url: config.vocab_url(), + local_name: config.vocab_file.filename().to_string(), + min_bytes: config.vocab_min_bytes, + }, + ] +} + +fn download_file( + url: &str, + local_name: &str, + min_bytes: u64, + model_dir: &Path, +) -> anyhow::Result<()> { + let local_path = model_dir.join(local_name); + let tmp_path = model_dir.join(format!("{local_name}.tmp")); + + tracing::info!("Downloading {local_name} ..."); + + let agent = ureq::AgentBuilder::new() + .timeout_connect(Duration::from_secs(30)) + .timeout_read(Duration::from_secs(300)) + .build(); + + let response = agent + .get(url) + .set("User-Agent", USER_AGENT) + .call() + .map_err(|e| anyhow::anyhow!("Failed to download {url}: {e}"))?; + + let status = response.status(); + if status != 200 { + anyhow::bail!("Download of {local_name} returned HTTP {status}"); + } + + let mut body = response.into_reader(); + let mut out = std::fs::File::create(&tmp_path)?; + let mut buf = vec![0u8; 65536]; + let mut total: u64 = 0; + + loop { + let n = body.read(&mut buf)?; + if n == 0 { + break; + } + std::io::Write::write_all(&mut out, &buf[..n])?; + total += n as u64; + } + drop(out); + + if total < min_bytes { + let _ = std::fs::remove_file(&tmp_path); + anyhow::bail!("Downloaded {local_name} is too small ({total} bytes, expected >= {min_bytes})"); + } + + std::fs::rename(&tmp_path, &local_path)?; + tracing::info!(" {local_name} — {:.1}MB saved", total as f64 / 1_048_576.0); + Ok(()) +} + +fn sha256_file(path: &Path) -> anyhow::Result { + use sha2::{Digest, Sha256}; + let mut file = std::fs::File::open(path)?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 65536]; + loop { + let n = file.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + let result = hasher.finalize(); + Ok(format!("{result:x}")) +} + +fn read_lockfile(model_dir: &Path) -> BTreeMap { + std::fs::read_to_string(model_dir.join(LOCKFILE)) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default() +} + +fn write_lockfile(model_dir: &Path, lock: &BTreeMap) -> anyhow::Result<()> { + if lock.is_empty() { + return Ok(()); + } + let json = serde_json::to_string_pretty(lock)?; + std::fs::write(model_dir.join(LOCKFILE), json)?; + Ok(()) +} + +pub fn clean_model(model_dir: &Path) -> anyhow::Result<()> { + for name in ["model.onnx", "vocab.txt", "tokenizer.json", LOCKFILE] { + let path = model_dir.join(name); + if path.exists() { + std::fs::remove_file(&path)?; + } + let tmp_path = model_dir.join(format!("{name}.tmp")); + if tmp_path.exists() { + std::fs::remove_file(&tmp_path)?; + } + } + Ok(()) +} diff --git a/crates/agentflare-store/src/embedding_pipeline/mod.rs b/crates/agentflare-store/src/embedding_pipeline/mod.rs new file mode 100644 index 00000000..026dab29 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/mod.rs @@ -0,0 +1,297 @@ +pub mod model_registry; +pub mod download; +pub mod tokenizer; +pub mod pooling; + +use std::path::{Path, PathBuf}; + +use model_registry::{EmbeddingModel, ModelConfig, VocabSource}; +use tokenizer::{TokenizedInput, WordPieceTokenizer, HfTokenizerWrapper}; + +pub struct EmbeddingEngine { + tokenizer: TokenizerKind, + dimensions: usize, + max_seq_len: usize, + model_id: EmbeddingModel, + model_config: ModelConfig, + session: std::sync::Mutex, + input_names: InputNames, + output_name: String, +} + +enum TokenizerKind { + WordPiece(WordPieceTokenizer), + HfTokenizer(HfTokenizerWrapper), +} + +struct InputNames { + input_ids: String, + attention_mask: String, + token_type_ids: Option, +} + +fn is_bert_punctuation(ch: char) -> bool { + if ch.is_ascii() { + matches!( + ch, + '!' | '"' | '#' | '$' | '%' | '&' | '\'' | '(' | ')' + | '*' | '+' | ',' | '-' | '.' | '/' | ':' | ';' + | '<' | '=' | '>' | '?' | '@' | '[' | '\\' | ']' + | '^' | '_' | '`' | '{' | '|' | '}' | '~' + ) + } else { + ch.is_ascii_punctuation() + } +} + +impl EmbeddingEngine { + pub fn load(model_dir: &Path) -> anyhow::Result { + let selected = model_registry::resolve_model(); + Self::load_model(model_dir, selected) + } + + pub fn load_model(base_dir: &Path, model_id: EmbeddingModel) -> anyhow::Result { + let config = model_id.config(); + let model_dir = base_dir.join(model_id.storage_dir_name()); + + download::ensure_model(&model_dir, &config)?; + + let tokenizer = load_tokenizer(&model_dir, &config)?; + let model_path = model_dir.join("model.onnx"); + + let mut session = ort::session::Session::builder() + .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))? + .with_intra_threads( + std::thread::available_parallelism().map_or(4, |n| n.get().max(1)), + ) + .map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))? + .with_optimization_level(ort::session::builder::GraphOptimizationLevel::All) + .map_err(|e| anyhow::anyhow!("ORT optimization: {e}"))? + .commit_from_file(&model_path) + .map_err(|e| anyhow::anyhow!("ORT load model: {e}"))?; + + let input_names_list: Vec = session + .inputs() + .iter() + .map(|i| i.name().to_string()) + .collect(); + + let token_type_ids = if config.needs_token_type_ids { + if input_names_list.len() < 3 { + anyhow::bail!( + "Model {} requires token_type_ids but only has {} inputs", + config.name, + input_names_list.len() + ); + } + Some(input_names_list[2].clone()) + } else if input_names_list.len() >= 3 { + Some(input_names_list[2].clone()) + } else { + None + }; + + let output_name = session + .outputs() + .first() + .map(|o| o.name().to_string()) + .ok_or_else(|| anyhow::anyhow!("Model has no named outputs"))?; + + let dimensions = detect_dimensions( + &config, + &tokenizer, + &input_names_list, + &token_type_ids, + &output_name, + )?; + + Ok(Self { + session: std::sync::Mutex::new(session), + tokenizer, + dimensions, + max_seq_len: config.max_seq_len, + model_id, + model_config: config, + input_names: InputNames { + input_ids: input_names_list[0].clone(), + attention_mask: input_names_list[1].clone(), + token_type_ids, + }, + output_name, + }) + } + + pub fn load_default() -> anyhow::Result { + Self::load(&Self::model_directory()) + } + + pub fn embed(&self, text: &str) -> anyhow::Result> { + let prefixed; + let input_text = if let Some(prefix) = &self.model_config.document_prefix { + prefixed = format!("{prefix}{text}"); + &prefixed + } else { + text + }; + let input = tokenize(&self.tokenizer, input_text, self.max_seq_len); + let mut hidden = self.run_inference(&input)?; + + let pooled = if let InputNames { token_type_ids: Some(_), .. } = &self.input_names { + pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions) + } else { + pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions) + }; + hidden = pooled; + + pooling::normalize_l2(&mut hidden); + Ok(hidden) + } + + pub fn embed_query(&self, query: &str) -> anyhow::Result> { + let prefixed; + let input_text = if let Some(prefix) = &self.model_config.query_prefix { + prefixed = format!("{prefix}{query}"); + &prefixed + } else { + query + }; + let input = tokenize(&self.tokenizer, input_text, self.max_seq_len); + let mut hidden = self.run_inference(&input)?; + + let pooled = pooling::mean_pool( + &hidden, + &input.attention_mask, + input.input_ids.len(), + self.dimensions, + ); + hidden = pooled; + + pooling::normalize_l2(&mut hidden); + Ok(hidden) + } + + pub fn dimensions(&self) -> usize { + self.dimensions + } + + pub fn model_name(&self) -> &str { + &self.model_config.name + } + + pub fn model_directory() -> PathBuf { + if let Ok(dir) = std::env::var("AGENTFLARE_MODELS_DIR") { + return PathBuf::from(dir); + } + let base = dirs::cache_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("agentflare") + .join("models"); + base + } + + fn run_inference(&self, input: &TokenizedInput) -> anyhow::Result> { + let seq_len = input.input_ids.len(); + let ids_vec: Vec = input.input_ids.iter().map(|&x| x as i64).collect(); + let mask_vec: Vec = input.attention_mask.iter().map(|&x| x as i64).collect(); + let ids_array = ndarray::Array2::from_shape_vec((1, seq_len), ids_vec)?; + let mask_array = ndarray::Array2::from_shape_vec((1, seq_len), mask_vec)?; + let ids_tensor = ort::value::Tensor::from_array(ids_array)?; + let mask_tensor = ort::value::Tensor::from_array(mask_array)?; + + let mut session = self + .session + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + if let Some(type_id) = &self.input_names.token_type_ids { + let type_vec: Vec = input.token_type_ids.iter().map(|&x| x as i64).collect(); + let type_array = ndarray::Array2::from_shape_vec((1, seq_len), type_vec)?; + let type_tensor = ort::value::Tensor::from_array(type_array)?; + let outputs = session.run(ort::inputs![ + self.input_names.input_ids.as_str() => ids_tensor, + self.input_names.attention_mask.as_str() => mask_tensor, + type_id.as_str() => type_tensor, + ])?; + let (_, data) = outputs[self.output_name.as_str()].try_extract_tensor::()?; + Ok(data.to_vec()) + } else { + let outputs = session.run(ort::inputs![ + self.input_names.input_ids.as_str() => ids_tensor, + self.input_names.attention_mask.as_str() => mask_tensor, + ])?; + let (_, data) = outputs[self.output_name.as_str()].try_extract_tensor::()?; + Ok(data.to_vec()) + } + } +} + +fn load_tokenizer(model_dir: &Path, config: &ModelConfig) -> anyhow::Result { + match &config.vocab_file { + VocabSource::VocabTxt(filename) => { + let path = model_dir.join(filename); + let tok = WordPieceTokenizer::from_file(&path)?; + Ok(TokenizerKind::WordPiece(tok)) + } + VocabSource::TokenizerJson(filename) => { + let path = model_dir.join(filename); + let tok = HfTokenizerWrapper::from_file(&path)?; + Ok(TokenizerKind::HfTokenizer(tok)) + } + } +} + +fn tokenize(tokenizer: &TokenizerKind, text: &str, max_len: usize) -> TokenizedInput { + match tokenizer { + TokenizerKind::WordPiece(wp) => wp.encode(text, max_len), + TokenizerKind::HfTokenizer(hf) => hf.encode(text, max_len), + } +} + +fn detect_dimensions( + config: &ModelConfig, + tokenizer: &TokenizerKind, + input_names: &[String], + token_type_ids: &Option, + output_name: &str, +) -> anyhow::Result { + let dummy = tokenize(tokenizer, "test", 8); + let seq_len = dummy.input_ids.len(); + if seq_len == 0 { + return Ok(config.dimensions); + } + + let ids_vec: Vec = dummy.input_ids.iter().map(|&x| x as i64).collect(); + let mask_vec: Vec = dummy.attention_mask.iter().map(|&x| x as i64).collect(); + let ids_array = ndarray::Array2::from_shape_vec((1, seq_len), ids_vec)?; + let mask_array = ndarray::Array2::from_shape_vec((1, seq_len), mask_vec)?; + let ids_tensor = ort::value::Tensor::from_array(ids_array)?; + let mask_tensor = ort::value::Tensor::from_array(mask_array)?; + + let mut session = ort::session::Session::builder() + .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))? + .with_intra_threads(1) + .map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))? + .commit_from_file(&config.model_url()) + .map_err(|_| anyhow::anyhow!("cannot probe dimensions without model file"))?; + + let outputs = if let Some(type_id) = token_type_ids { + let type_vec: Vec = dummy.token_type_ids.iter().map(|&x| x as i64).collect(); + let type_array = ndarray::Array2::from_shape_vec((1, seq_len), type_vec)?; + let type_tensor = ort::value::Tensor::from_array(type_array)?; + session.run(ort::inputs![ + input_names[0].as_str() => ids_tensor, + input_names[1].as_str() => mask_tensor, + type_id.as_str() => type_tensor, + ])? + } else { + session.run(ort::inputs![ + input_names[0].as_str() => ids_tensor, + input_names[1].as_str() => mask_tensor, + ])? + }; + + let (shape, _) = outputs[output_name].try_extract_tensor::()?; + shape.last().copied().map(|s| s as usize).ok_or_else(|| { + anyhow::anyhow!("could not detect embedding dimensions from model output") + }) +} diff --git a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs new file mode 100644 index 00000000..d9b57612 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs @@ -0,0 +1,200 @@ +use std::fmt; + +#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum EmbeddingModel { + AllMiniLmL6V2, + NomicEmbedV1_5, + Custom(CustomModelSpec), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub struct CustomModelSpec { + pub repo: String, + pub revision: Option, + pub dimensions: Option, +} + +impl CustomModelSpec { + fn parse(s: &str) -> Option { + let (repo, revision) = match s.split_once('@') { + Some((r, rev)) => ( + r.trim(), + Some(rev.trim().to_string()).filter(|v| !v.is_empty()), + ), + None => (s.trim(), None), + }; + let mut parts = repo.split('/'); + let (owner, name) = (parts.next()?, parts.next()?); + if parts.next().is_some() || owner.is_empty() || name.is_empty() + || repo.chars().any(char::is_whitespace) + { + return None; + } + Some(Self { repo: repo.to_string(), revision, dimensions: None }) + } + + fn storage_slug(&self) -> String { + let mut slug = String::from("hf-"); + for c in self.repo.chars() { + slug.push(match c { + 'a'..='z' | '0'..='9' | '-' => c, + 'A'..='Z' => c.to_ascii_lowercase(), + _ => '-', + }); + } + if let Some(rev) = &self.revision { + slug.push('-'); + for c in rev.chars().take(16) { + slug.push(match c { + 'a'..='z' | '0'..='9' | '-' => c, + 'A'..='Z' => c.to_ascii_lowercase(), + _ => '-', + }); + } + } + slug + } +} + +impl EmbeddingModel { + pub const DEFAULT: Self = Self::AllMiniLmL6V2; + + pub fn config(&self) -> ModelConfig { + match self { + Self::AllMiniLmL6V2 => ModelConfig { + model: self.clone(), + name: "all-MiniLM-L6-v2".into(), + hf_repo: "sentence-transformers/all-MiniLM-L6-v2".into(), + revision: None, + onnx_path: "onnx/model.onnx".into(), + vocab_file: VocabSource::VocabTxt("vocab.txt".into()), + dimensions: 384, + max_seq_len: 256, + model_min_bytes: 1_000_000, + vocab_min_bytes: 100_000, + query_prefix: None, + document_prefix: None, + needs_token_type_ids: true, + }, + Self::NomicEmbedV1_5 => ModelConfig { + model: self.clone(), + name: "nomic-embed-text-v1.5".into(), + hf_repo: "nomic-ai/nomic-embed-text-v1.5".into(), + revision: None, + onnx_path: "onnx/model.onnx".into(), + vocab_file: VocabSource::VocabTxt("vocab.txt".into()), + dimensions: 768, + max_seq_len: 512, + model_min_bytes: 100_000_000, + vocab_min_bytes: 100_000, + query_prefix: Some("search_query: ".into()), + document_prefix: Some("search_document: ".into()), + needs_token_type_ids: false, + }, + Self::Custom(spec) => ModelConfig { + model: self.clone(), + name: match &spec.revision { + Some(rev) => format!("hf:{}@{rev}", spec.repo), + None => format!("hf:{}", spec.repo), + }, + hf_repo: spec.repo.clone(), + revision: spec.revision.clone(), + onnx_path: "onnx/model.onnx".into(), + vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()), + dimensions: spec.dimensions.unwrap_or(768), + max_seq_len: 512, + model_min_bytes: 1_000_000, + vocab_min_bytes: 1_000, + query_prefix: None, + document_prefix: None, + needs_token_type_ids: false, + }, + } + } + + pub fn from_str_name(s: &str) -> Option { + let trimmed = s.trim(); + if let Some(rest) = trimmed.strip_prefix("hf:") { + return CustomModelSpec::parse(rest).map(Self::Custom); + } + match trimmed.to_lowercase().replace('_', "-").as_str() { + "all-minilm-l6-v2" | "minilm" | "default" => Some(Self::AllMiniLmL6V2), + "nomic-embed-v1.5" | "nomic-embed-text-v1.5" | "nomic" | "nomic-embed" => { + Some(Self::NomicEmbedV1_5) + } + _ => None, + } + } + + pub const ALL: &'static [Self] = &[Self::AllMiniLmL6V2, Self::NomicEmbedV1_5]; + + pub fn storage_dir_name(&self) -> String { + match self { + Self::AllMiniLmL6V2 => "all-minilm-l6-v2".to_string(), + Self::NomicEmbedV1_5 => "nomic-embed-v1.5".to_string(), + Self::Custom(spec) => spec.storage_slug(), + } + } +} + +impl fmt::Display for EmbeddingModel { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.config().name) + } +} + +#[derive(Debug, Clone)] +pub enum VocabSource { + VocabTxt(String), + TokenizerJson(String), +} + +impl VocabSource { + pub fn filename(&self) -> &str { + match self { + Self::VocabTxt(f) | Self::TokenizerJson(f) => f, + } + } +} + +#[derive(Debug, Clone)] +pub struct ModelConfig { + pub model: EmbeddingModel, + pub name: String, + pub hf_repo: String, + pub revision: Option, + pub onnx_path: String, + pub vocab_file: VocabSource, + pub dimensions: usize, + pub max_seq_len: usize, + pub model_min_bytes: u64, + pub vocab_min_bytes: u64, + pub query_prefix: Option, + pub document_prefix: Option, + pub needs_token_type_ids: bool, +} + +impl ModelConfig { + fn resolve_base(&self) -> String { + format!( + "https://huggingface.co/{}/resolve/{}", + self.hf_repo, + self.revision.as_deref().unwrap_or("main") + ) + } + + pub fn model_url(&self) -> String { + format!("{}/{}", self.resolve_base(), self.onnx_path) + } + + pub fn vocab_url(&self) -> String { + format!("{}/{}", self.resolve_base(), self.vocab_file.filename()) + } +} + +pub fn resolve_model() -> EmbeddingModel { + let env_val = std::env::var("AGENTFLARE_EMBEDDING_MODEL").ok(); + let name = env_val.as_deref().unwrap_or("default"); + EmbeddingModel::from_str_name(name).unwrap_or(EmbeddingModel::DEFAULT) +} diff --git a/crates/agentflare-store/src/embedding_pipeline/pooling.rs b/crates/agentflare-store/src/embedding_pipeline/pooling.rs new file mode 100644 index 00000000..a7bf46e5 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/pooling.rs @@ -0,0 +1,41 @@ +pub fn mean_pool( + hidden_states: &[f32], + attention_mask: &[i32], + seq_len: usize, + dim: usize, +) -> Vec { + let mut sum = vec![0.0f32; dim]; + let mut count = 0.0f32; + + for pos in 0..seq_len { + if attention_mask.get(pos).copied().unwrap_or(0) > 0 { + let offset = pos * dim; + for d in 0..dim { + if let Some(&val) = hidden_states.get(offset + d) { + sum[d] += val; + } + } + count += 1.0; + } + } + + if count > 0.0 { + for val in &mut sum { + *val /= count; + } + } + sum +} + +pub fn normalize_l2(vec: &mut [f32]) { + let norm: f32 = vec.iter().map(|x| x * x).sum::().sqrt(); + if norm > f32::EPSILON { + for x in vec.iter_mut() { + *x /= norm; + } + } +} + +pub fn l2_norm(vec: &[f32]) -> f32 { + vec.iter().map(|x| x * x).sum::().sqrt() +} diff --git a/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs new file mode 100644 index 00000000..ea094382 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs @@ -0,0 +1,204 @@ +use std::collections::HashMap; +use std::path::Path; + +pub struct WordPieceTokenizer { + vocab: HashMap, + pub cls_id: i32, + pub sep_id: i32, + pub pad_id: i32, + pub unk_id: i32, + max_word_chars: usize, +} + +#[derive(Debug, Clone)] +pub struct TokenizedInput { + pub input_ids: Vec, + pub attention_mask: Vec, + pub token_type_ids: Vec, +} + +impl WordPieceTokenizer { + pub fn from_file(path: &Path) -> anyhow::Result { + let content = std::fs::read_to_string(path) + .map_err(|e| anyhow::anyhow!("Failed to read vocab file {}: {}", path.display(), e))?; + Self::from_vocab_str(&content) + } + + pub fn from_vocab_str(vocab_str: &str) -> anyhow::Result { + let vocab: HashMap = vocab_str + .lines() + .enumerate() + .map(|(i, line)| (line.to_string(), i as i32)) + .collect(); + + let cls_id = *vocab.get("[CLS]") + .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [CLS] token"))?; + let sep_id = *vocab.get("[SEP]") + .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [SEP] token"))?; + let pad_id = *vocab.get("[PAD]") + .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [PAD] token"))?; + let unk_id = *vocab.get("[UNK]") + .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [UNK] token"))?; + + Ok(Self { vocab, cls_id, sep_id, pad_id, unk_id, max_word_chars: 200 }) + } + + pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { + let words = self.pre_tokenize(text); + let mut ids = vec![self.cls_id]; + + for word in &words { + if ids.len() >= max_len - 1 { break; } + let subword_ids = self.wordpiece_encode(word); + for id in subword_ids { + if ids.len() >= max_len - 1 { break; } + ids.push(id); + } + } + ids.push(self.sep_id); + + let len = ids.len(); + TokenizedInput { + input_ids: ids, + attention_mask: vec![1; len], + token_type_ids: vec![0; len], + } + } + + fn pre_tokenize(&self, text: &str) -> Vec { + let mut words = Vec::new(); + let mut current = String::new(); + + for ch in text.chars() { + if ch.is_whitespace() { + if !current.is_empty() { + words.extend(self.split_identifier(¤t)); + current.clear(); + } + } else if is_bert_punctuation(ch) { + if !current.is_empty() { + words.extend(self.split_identifier(¤t)); + current.clear(); + } + words.push(ch.to_string()); + } else { + current.push(ch); + } + } + if !current.is_empty() { + words.extend(self.split_identifier(¤t)); + } + words.iter().map(|w| w.to_lowercase()).collect() + } + + fn split_identifier(&self, word: &str) -> Vec { + let lower = word.to_lowercase(); + if self.vocab.contains_key(&lower) { + return vec![word.to_string()]; + } + let mut parts = Vec::new(); + let mut current = String::new(); + let chars: Vec = word.chars().collect(); + for (i, &ch) in chars.iter().enumerate() { + if ch == '_' || ch == '-' { + if !current.is_empty() { parts.push(current.clone()); current.clear(); } + } else if i > 0 && ch.is_ascii_uppercase() && chars[i - 1].is_ascii_lowercase() { + if !current.is_empty() { parts.push(current.clone()); current.clear(); } + current.push(ch); + } else { + current.push(ch); + } + } + if !current.is_empty() { parts.push(current); } + if parts.is_empty() { vec![word.to_string()] } else { parts } + } + + fn wordpiece_encode(&self, word: &str) -> Vec { + if word.chars().count() > self.max_word_chars { + return vec![self.unk_id]; + } + let chars: Vec = word.chars().collect(); + let mut tokens = Vec::new(); + let mut start = 0; + while start < chars.len() { + let mut end = chars.len(); + let mut matched = false; + while start < end { + let substr: String = chars[start..end].iter().collect(); + let candidate = if start > 0 { format!("##{substr}") } else { substr }; + if let Some(&id) = self.vocab.get(&candidate) { + tokens.push(id); + matched = true; + start = end; + break; + } + end -= 1; + } + if !matched { + tokens.push(self.unk_id); + start += 1; + } + } + tokens + } +} + +pub struct HfTokenizerWrapper { + inner: WordPieceTokenizer, +} + +impl HfTokenizerWrapper { + pub fn from_file(path: &Path) -> anyhow::Result { + let content = std::fs::read_to_string(path)?; + Self::from_json(&content) + } + + fn from_json(json_str: &str) -> anyhow::Result { + let parsed: serde_json::Value = serde_json::from_str(json_str)?; + let vocab_obj = parsed + .get("model") + .and_then(|m| m.get("vocab")) + .and_then(|v| v.as_object()) + .ok_or_else(|| anyhow::anyhow!("tokenizer.json missing model.vocab object"))?; + + let mut vocab_lines: Vec<(String, i32)> = vocab_obj + .iter() + .filter_map(|(token, id)| id.as_i64().map(|id| (token.clone(), id as i32))) + .collect(); + vocab_lines.sort_by_key(|(_, id)| *id); + + for (token, _) in &mut vocab_lines { + let mapped = match token.as_str() { + "" => "[CLS]", "" => "[SEP]", + "" => "[PAD]", "" => "[UNK]", + "" => "[MASK]", _ => continue, + }; + *token = mapped.to_string(); + } + + let vocab_str: String = vocab_lines.into_iter() + .map(|(token, _)| token) + .collect::>() + .join("\n"); + + let inner = WordPieceTokenizer::from_vocab_str(&vocab_str)?; + Ok(Self { inner }) + } + + pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { + self.inner.encode(text, max_len) + } +} + +fn is_bert_punctuation(ch: char) -> bool { + if ch.is_ascii() { + matches!(ch, + '!' | '"' | '#' | '$' | '%' | '&' | '\'' | '(' | ')' + | '*' | '+' | ',' | '-' | '.' | '/' | ':' | ';' + | '<' | '=' | '>' | '?' | '@' | '[' | '\\' | ']' + | '^' | '_' | '`' | '{' | '|' | '}' | '~' + ) + } else { + ch.is_ascii_punctuation() + } +} diff --git a/crates/agentflare-store/src/kv.rs b/crates/agentflare-store/src/kv.rs index ea1c77e3..bb176ccd 100644 --- a/crates/agentflare-store/src/kv.rs +++ b/crates/agentflare-store/src/kv.rs @@ -13,8 +13,9 @@ pub struct KvEntry { impl Store { pub fn kv_set(&self, key: &str, value: &[u8]) -> rusqlite::Result<()> { + let conn = self.conn(); let now = db_kit::ids::now(); - self.conn.execute( + conn.execute( "INSERT INTO store_kv (key, value, created_at, updated_at) VALUES (?1, ?2, ?3, ?3) ON CONFLICT(key) DO UPDATE SET value = ?2, updated_at = ?3", @@ -24,8 +25,8 @@ impl Store { } pub fn kv_get(&self, key: &str) -> rusqlite::Result> { - self.conn - .query_row( + let conn = self.conn(); + conn.query_row( "SELECT key, value, created_at, updated_at FROM store_kv WHERE key = ?1", params![key], |row| { @@ -41,8 +42,8 @@ impl Store { } pub fn kv_exists(&self, key: &str) -> rusqlite::Result { - self.conn - .query_row( + let conn = self.conn(); + conn.query_row( "SELECT 1 FROM store_kv WHERE key = ?1", params![key], |_| Ok(()), @@ -52,16 +53,14 @@ impl Store { } pub fn kv_delete(&self, key: &str) -> rusqlite::Result { - let n = self - .conn - .execute("DELETE FROM store_kv WHERE key = ?1", params![key])?; + let conn = self.conn(); + let n = conn.execute("DELETE FROM store_kv WHERE key = ?1", params![key])?; Ok(n > 0) } pub fn kv_scan(&self, prefix: &str) -> rusqlite::Result> { - let mut stmt = self - .conn - .prepare("SELECT key, value, created_at, updated_at FROM store_kv WHERE key >= ?1 AND key < ?2 ORDER BY key")?; + let conn = self.conn(); + let mut stmt = conn.prepare("SELECT key, value, created_at, updated_at FROM store_kv WHERE key >= ?1 AND key < ?2 ORDER BY key")?; let end = { let mut s = prefix.to_string(); s.push('\u{10FFFF}'); diff --git a/crates/agentflare-store/src/lib.rs b/crates/agentflare-store/src/lib.rs index 41d4d8e8..61392244 100644 --- a/crates/agentflare-store/src/lib.rs +++ b/crates/agentflare-store/src/lib.rs @@ -6,6 +6,9 @@ pub mod leases; pub mod migrate; pub mod migrations; +#[cfg(feature = "embeddings")] +pub mod embedding_pipeline; + use rusqlite::Connection; use std::path::{Path, PathBuf}; @@ -26,7 +29,7 @@ pub enum Error { } pub struct Store { - conn: Connection, + conn: parking_lot::Mutex, root: PathBuf, } @@ -34,19 +37,22 @@ impl Store { pub fn open_file(path: &Path) -> Result { let conn = db_kit::open_file(path, &migrations::migrations())?; let root = path.parent().unwrap_or(Path::new(".")).to_path_buf(); - Ok(Self { conn, root }) + Ok(Self { + conn: parking_lot::Mutex::new(conn), + root, + }) } pub fn open_memory() -> Result { let conn = db_kit::open_memory(&migrations::migrations())?; Ok(Self { - conn, + conn: parking_lot::Mutex::new(conn), root: PathBuf::from(":memory:"), }) } - pub fn conn(&self) -> &Connection { - &self.conn + pub fn conn(&self) -> parking_lot::MutexGuard<'_, Connection> { + self.conn.lock() } pub fn root(&self) -> &Path { diff --git a/crates/agentflare-store/src/migrations.rs b/crates/agentflare-store/src/migrations.rs index 67b3e17a..032b6334 100644 --- a/crates/agentflare-store/src/migrations.rs +++ b/crates/agentflare-store/src/migrations.rs @@ -57,5 +57,28 @@ pub fn migrations() -> Migrations<'static> { heartbeat_at INTEGER NOT NULL );", ), + M::up( + "ALTER TABLE store_documents ADD COLUMN title TEXT NOT NULL DEFAULT ''; + ALTER TABLE store_documents ADD COLUMN doc_type TEXT NOT NULL DEFAULT 'file'; + ALTER TABLE store_documents ADD COLUMN blob_hash TEXT; + ALTER TABLE store_documents ADD COLUMN mime TEXT NOT NULL DEFAULT ''; + ALTER TABLE store_documents ADD COLUMN tags TEXT NOT NULL DEFAULT '[]'; + ALTER TABLE store_documents ADD COLUMN session_id TEXT; + ALTER TABLE store_documents ADD COLUMN source TEXT NOT NULL DEFAULT ''; + ALTER TABLE store_documents ADD COLUMN version INTEGER NOT NULL DEFAULT 1;", + ), + M::up( + "CREATE TABLE IF NOT EXISTS store_doc_history ( + id TEXT PRIMARY KEY NOT NULL, + doc_id TEXT NOT NULL REFERENCES store_documents(id), + version INTEGER NOT NULL, + content TEXT NOT NULL DEFAULT '', + blob_hash TEXT, + mime TEXT NOT NULL DEFAULT '', + title TEXT NOT NULL DEFAULT '', + created_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_doc_history_doc ON store_doc_history(doc_id);", + ), ]) } diff --git a/src/main.rs b/src/main.rs index be0b9619..03a049bd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -41,6 +41,7 @@ mod rollup; mod rule_text; mod shell; mod state; +mod store; mod tool_install; mod ui; mod uninstall; diff --git a/src/store.rs b/src/store.rs new file mode 100644 index 00000000..d4cbaf41 --- /dev/null +++ b/src/store.rs @@ -0,0 +1,36 @@ +use agentflare_store::Store; +use std::path::PathBuf; +use std::sync::OnceLock; + +static STORE: OnceLock> = OnceLock::new(); + +pub fn store_path() -> PathBuf { + crate::paths::home().join(".agentflare").join("agentflare.db") +} + +pub fn open() -> &'static Store { + STORE + .get_or_init(|| { + let path = store_path(); + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + Store::open_file(&path) + }) + .as_ref() + .expect("failed to open agentflare store") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::paths::test_support::with_temp_home; + + #[test] + fn open_and_close() { + with_temp_home(|| { + let store = open(); + store.conn().execute_batch("SELECT 1").unwrap(); + }); + } +} From 6f9d4e0ceeeabd17710f8fe83c16616b1db8b3a7 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 15:02:14 +0530 Subject: [PATCH 05/16] fix(agentflare-store): wrap doc upsert in a transaction, preserve blob_hash/mime in history, fix broken embedding dimension probe - doc_upsert_with_opts issued ~8 unwrapped statements against a WAL-mode, multi-process-shared SQLite file; wrap the whole upsert in a transaction so concurrent readers/crashes can't observe a half-written document. - store_doc_history never captured blob_hash/mime, so doc_history()/ doc_get_version() always returned blank metadata for old versions; capture and store them, add a regression test. - doc_vec_search silently swallowed row-decode errors via filter_map(Result::ok), making a real DB error look like zero matches; propagate them instead. - detect_dimensions() passed a remote https:// model URL into ort's commit_from_file (which expects a local path), so EmbeddingEngine::load() always failed when the embeddings feature was used; pass the local model path that load_model() already resolves. - removed dead .meta sidecar file writes/deletes in blobs.rs (never read), a duplicate is_bert_punctuation() in embedding_pipeline/mod.rs, a no-op if/else in embed() with identical branches, and two warnings (unused import, unused mut). --- crates/agentflare-store/src/blobs.rs | 2 - crates/agentflare-store/src/documents.rs | 76 ++++++++++++++----- .../src/embedding_pipeline/download.rs | 2 +- .../src/embedding_pipeline/mod.rs | 27 ++----- 4 files changed, 65 insertions(+), 42 deletions(-) diff --git a/crates/agentflare-store/src/blobs.rs b/crates/agentflare-store/src/blobs.rs index fe1c81c1..c89aed51 100644 --- a/crates/agentflare-store/src/blobs.rs +++ b/crates/agentflare-store/src/blobs.rs @@ -28,14 +28,12 @@ fn write_disk_blob(root: &Path, hash: &str, data: &[u8]) -> Result<(), std::io:: std::fs::create_dir_all(parent)?; } std::fs::write(&path, data)?; - std::fs::write(path.with_extension("meta"), &[])?; Ok(()) } fn delete_disk_blob(root: &Path, hash: &str) { let path = blob_disk_path(root, hash); let _ = std::fs::remove_file(&path); - let _ = std::fs::remove_file(path.with_extension("meta")); } use std::path::PathBuf; diff --git a/crates/agentflare-store/src/documents.rs b/crates/agentflare-store/src/documents.rs index 9e141732..aaa4c9b3 100644 --- a/crates/agentflare-store/src/documents.rs +++ b/crates/agentflare-store/src/documents.rs @@ -109,7 +109,7 @@ impl Store { let existing = conn .query_row( - "SELECT id, rowid, content, version FROM store_documents + "SELECT id, rowid, content, version, blob_hash, mime FROM store_documents WHERE project_id = ?1 AND path = ?2", params![project_id, path], |row| { @@ -118,23 +118,27 @@ impl Store { row.get::<_, i64>(1)?, row.get::<_, String>(2)?, row.get::<_, i32>(3)?, + row.get::<_, Option>(4)?, + row.get::<_, String>(5)?, )) }, ) .optional()?; - if let Some((existing_id, rowid, old_content, old_version)) = existing { + if let Some((existing_id, rowid, old_content, old_version, old_blob_hash, old_mime)) = existing { let new_version = old_version + 1; let history_id = db_kit::ids::new_id(); + let tx = conn.unchecked_transaction()?; + // Snapshot current version to history - conn.execute( - "INSERT INTO store_doc_history (id, doc_id, version, content, title, created_at) - VALUES (?1, ?2, ?3, ?4, (SELECT title FROM store_documents WHERE id = ?2), ?5)", - params![history_id, existing_id, old_version, old_content, now], + tx.execute( + "INSERT INTO store_doc_history (id, doc_id, version, content, blob_hash, mime, title, created_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, (SELECT title FROM store_documents WHERE id = ?2), ?7)", + params![history_id, existing_id, old_version, old_content, old_blob_hash, old_mime, now], )?; - conn.execute( + tx.execute( "UPDATE store_documents SET content = ?1, updated_at = ?2, deleted_at = NULL, version = ?3 @@ -144,50 +148,51 @@ impl Store { // Apply optional updates (need separate UPDATE to avoid long SQL) if let Some(title) = &opts.title { - conn.execute( + tx.execute( "UPDATE store_documents SET title = ?1 WHERE id = ?2", params![title, existing_id], )?; } if let Some(doc_type) = &opts.doc_type { - conn.execute( + tx.execute( "UPDATE store_documents SET doc_type = ?1 WHERE id = ?2", params![doc_type, existing_id], )?; } if opts.blob_hash.is_some() { - conn.execute( + tx.execute( "UPDATE store_documents SET blob_hash = ?1 WHERE id = ?2", params![opts.blob_hash, existing_id], )?; } if let Some(mime) = &opts.mime { - conn.execute( + tx.execute( "UPDATE store_documents SET mime = ?1 WHERE id = ?2", params![mime, existing_id], )?; } if let Some(tags) = &opts.tags { let json = serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string()); - conn.execute( + tx.execute( "UPDATE store_documents SET tags = ?1 WHERE id = ?2", params![json, existing_id], )?; } if opts.session_id.is_some() { - conn.execute( + tx.execute( "UPDATE store_documents SET session_id = ?1 WHERE id = ?2", params![opts.session_id, existing_id], )?; } if let Some(source) = &opts.source { - conn.execute( + tx.execute( "UPDATE store_documents SET source = ?1 WHERE id = ?2", params![source, existing_id], )?; } - Self::doc_sync_fts(&conn, rowid, content)?; + Self::doc_sync_fts(&tx, rowid, content)?; + tx.commit()?; drop(conn); self.doc_get(&existing_id).map(|o| o.unwrap()) } else { @@ -406,7 +411,7 @@ impl Store { JOIN store_documents d ON d.id = v.doc_id WHERE d.project_id = ?1 AND d.deleted_at IS NULL", )?; - let mut results: Vec<(f64, DocMatch)> = stmt + let rows: Vec<(String, String, String, Vec)> = stmt .query_map(params![project_id], |row| { let id: String = row.get(0)?; let project_id: String = row.get(1)?; @@ -414,7 +419,9 @@ impl Store { let blob: Vec = row.get(3)?; Ok((id, project_id, path, blob)) })? - .filter_map(|r| r.ok()) + .collect::>>()?; + let mut results: Vec<(f64, DocMatch)> = rows + .into_iter() .filter_map(|(id, pid, path, blob)| { if blob.len() % 4 != 0 { return None; @@ -660,4 +667,39 @@ mod tests { assert_eq!(history[0].version, 1); assert_eq!(history[0].content, "v1"); } + + #[test] + fn history_snapshot_preserves_blob_hash_and_mime() { + let s = store(); + let doc = s + .doc_upsert_with_opts( + "p", + "/v.md", + "v1", + DocUpsertOpts { + blob_hash: Some("hash-v1".into()), + mime: Some("text/plain".into()), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(doc.version, 1); + + s.doc_upsert_with_opts( + "p", + "/v.md", + "v2", + DocUpsertOpts { + blob_hash: Some("hash-v2".into()), + mime: Some("text/markdown".into()), + ..Default::default() + }, + ) + .unwrap(); + + let history = s.doc_history(&doc.id).unwrap(); + assert_eq!(history.len(), 1); + assert_eq!(history[0].blob_hash.as_deref(), Some("hash-v1")); + assert_eq!(history[0].mime, "text/plain"); + } } diff --git a/crates/agentflare-store/src/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs index bd4859fd..58b7aa3f 100644 --- a/crates/agentflare-store/src/embedding_pipeline/download.rs +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -3,7 +3,7 @@ use std::io::Read; use std::path::{Path, PathBuf}; use std::time::Duration; -use super::model_registry::{ModelConfig, VocabSource}; +use super::model_registry::ModelConfig; const USER_AGENT: &str = concat!("agentflare-store/", env!("CARGO_PKG_VERSION")); const LOCKFILE: &str = "model.lock.json"; diff --git a/crates/agentflare-store/src/embedding_pipeline/mod.rs b/crates/agentflare-store/src/embedding_pipeline/mod.rs index 026dab29..04ebee53 100644 --- a/crates/agentflare-store/src/embedding_pipeline/mod.rs +++ b/crates/agentflare-store/src/embedding_pipeline/mod.rs @@ -30,20 +30,6 @@ struct InputNames { token_type_ids: Option, } -fn is_bert_punctuation(ch: char) -> bool { - if ch.is_ascii() { - matches!( - ch, - '!' | '"' | '#' | '$' | '%' | '&' | '\'' | '(' | ')' - | '*' | '+' | ',' | '-' | '.' | '/' | ':' | ';' - | '<' | '=' | '>' | '?' | '@' | '[' | '\\' | ']' - | '^' | '_' | '`' | '{' | '|' | '}' | '~' - ) - } else { - ch.is_ascii_punctuation() - } -} - impl EmbeddingEngine { pub fn load(model_dir: &Path) -> anyhow::Result { let selected = model_registry::resolve_model(); @@ -59,7 +45,7 @@ impl EmbeddingEngine { let tokenizer = load_tokenizer(&model_dir, &config)?; let model_path = model_dir.join("model.onnx"); - let mut session = ort::session::Session::builder() + let session = ort::session::Session::builder() .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))? .with_intra_threads( std::thread::available_parallelism().map_or(4, |n| n.get().max(1)), @@ -99,6 +85,7 @@ impl EmbeddingEngine { let dimensions = detect_dimensions( &config, + &model_path, &tokenizer, &input_names_list, &token_type_ids, @@ -136,12 +123,7 @@ impl EmbeddingEngine { let input = tokenize(&self.tokenizer, input_text, self.max_seq_len); let mut hidden = self.run_inference(&input)?; - let pooled = if let InputNames { token_type_ids: Some(_), .. } = &self.input_names { - pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions) - } else { - pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions) - }; - hidden = pooled; + hidden = pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions); pooling::normalize_l2(&mut hidden); Ok(hidden) @@ -249,6 +231,7 @@ fn tokenize(tokenizer: &TokenizerKind, text: &str, max_len: usize) -> TokenizedI fn detect_dimensions( config: &ModelConfig, + model_path: &Path, tokenizer: &TokenizerKind, input_names: &[String], token_type_ids: &Option, @@ -271,7 +254,7 @@ fn detect_dimensions( .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))? .with_intra_threads(1) .map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))? - .commit_from_file(&config.model_url()) + .commit_from_file(model_path) .map_err(|_| anyhow::anyhow!("cannot probe dimensions without model file"))?; let outputs = if let Some(type_id) = token_type_ids { From bb4909e74d4db09b9fc9dbe9a42ab5883661a9c6 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 23:45:47 +0530 Subject: [PATCH 06/16] fix(agentflare-store): dedupe normalize, cargo fmt, clippy clean - embedding_pipeline::pooling::normalize_l2/l2_norm duplicated embed::normalize with a looser epsilon (f32::EPSILON vs 1e-12) and l2_norm had zero callers. Removed both; embed::normalize is now the single L2-normalize implementation, used by embed() and embed_query(). - cargo fmt across the crate (documents.rs, mod.rs, model_registry.rs, pooling.rs, kv.rs -- download.rs/tokenizer.rs left as-is, in-flight elsewhere). - clippy -D warnings clean: dead model_id field (added a model_id() accessor alongside the existing dimensions() getter), a manual index loop over sum flagged by clippy::needless_range_loop, and a let-and-return in model_directory(). --- crates/agentflare-store/src/documents.rs | 121 +++++++++++------- .../src/embedding_pipeline/mod.rs | 41 +++--- .../src/embedding_pipeline/model_registry.rs | 10 +- .../src/embedding_pipeline/pooling.rs | 17 +-- crates/agentflare-store/src/kv.rs | 36 +++--- 5 files changed, 126 insertions(+), 99 deletions(-) diff --git a/crates/agentflare-store/src/documents.rs b/crates/agentflare-store/src/documents.rs index aaa4c9b3..76188d38 100644 --- a/crates/agentflare-store/src/documents.rs +++ b/crates/agentflare-store/src/documents.rs @@ -54,7 +54,11 @@ pub struct DocUpsertOpts { } impl Store { - fn doc_sync_fts(conn: &rusqlite::Connection, row_id: i64, content: &str) -> rusqlite::Result<()> { + fn doc_sync_fts( + conn: &rusqlite::Connection, + row_id: i64, + content: &str, + ) -> rusqlite::Result<()> { conn.execute( "DELETE FROM store_docs_fts WHERE rowid = ?1", params![row_id], @@ -125,7 +129,9 @@ impl Store { ) .optional()?; - if let Some((existing_id, rowid, old_content, old_version, old_blob_hash, old_mime)) = existing { + if let Some((existing_id, rowid, old_content, old_version, old_blob_hash, old_mime)) = + existing + { let new_version = old_version + 1; let history_id = db_kit::ids::new_id(); @@ -238,13 +244,13 @@ impl Store { pub fn doc_get(&self, id: &str) -> rusqlite::Result> { let conn = self.conn(); conn.query_row( - "SELECT id, project_id, path, content, title, doc_type, blob_hash, mime, tags, + "SELECT id, project_id, path, content, title, doc_type, blob_hash, mime, tags, session_id, source, version, created_at, updated_at, deleted_at FROM store_documents WHERE id = ?1", - params![id], - Self::row_to_document, - ) - .optional() + params![id], + Self::row_to_document, + ) + .optional() } pub fn doc_delete(&self, id: &str) -> rusqlite::Result { @@ -316,26 +322,30 @@ impl Store { rows.collect() } - pub fn doc_get_version(&self, doc_id: &str, version: i32) -> rusqlite::Result> { + pub fn doc_get_version( + &self, + doc_id: &str, + version: i32, + ) -> rusqlite::Result> { let conn = self.conn(); conn.query_row( - "SELECT id, doc_id, version, content, blob_hash, mime, title, created_at + "SELECT id, doc_id, version, content, blob_hash, mime, title, created_at FROM store_doc_history WHERE doc_id = ?1 AND version = ?2", - params![doc_id, version], - |row| { - Ok(DocVersion { - id: row.get(0)?, - doc_id: row.get(1)?, - version: row.get(2)?, - content: row.get(3)?, - blob_hash: row.get(4)?, - mime: row.get(5)?, - title: row.get(6)?, - created_at: row.get(7)?, - }) - }, - ) - .optional() + params![doc_id, version], + |row| { + Ok(DocVersion { + id: row.get(0)?, + doc_id: row.get(1)?, + version: row.get(2)?, + content: row.get(3)?, + blob_hash: row.get(4)?, + mime: row.get(5)?, + title: row.get(6)?, + created_at: row.get(7)?, + }) + }, + ) + .optional() } pub fn doc_search( @@ -384,18 +394,18 @@ impl Store { pub fn doc_get_embedding(&self, doc_id: &str) -> rusqlite::Result>> { let conn = self.conn(); conn.query_row( - "SELECT embedding FROM store_docs_vec WHERE doc_id = ?1", - params![doc_id], - |row| { - let blob: Vec = row.get(0)?; - let vec: Vec = blob - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - Ok(vec) - }, - ) - .optional() + "SELECT embedding FROM store_docs_vec WHERE doc_id = ?1", + params![doc_id], + |row| { + let blob: Vec = row.get(0)?; + let vec: Vec = blob + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + Ok(vec) + }, + ) + .optional() } pub fn doc_vec_search( @@ -461,11 +471,19 @@ impl Store { let mut max_fts = fts.first().map(|m| m.score).unwrap_or(1.0); let mut max_vec = vec.first().map(|m| m.score).unwrap_or(1.0); - if max_fts < 1e-12 { max_fts = 1.0; } - if max_vec < 1e-12 { max_vec = 1.0; } + if max_fts < 1e-12 { + max_fts = 1.0; + } + if max_vec < 1e-12 { + max_vec = 1.0; + } - for m in &mut fts { m.score = alpha * (m.score / max_fts); } - for m in &mut vec { m.score = (1.0 - alpha) * (m.score / max_vec); } + for m in &mut fts { + m.score = alpha * (m.score / max_fts); + } + for m in &mut vec { + m.score = (1.0 - alpha) * (m.score / max_vec); + } let mut combined: Vec = Vec::new(); let mut seen = std::collections::HashSet::new(); @@ -637,14 +655,21 @@ mod tests { #[test] fn upsert_with_metadata() { let s = store(); - let doc = s.doc_upsert_with_opts("p", "/meta.md", "content", DocUpsertOpts { - title: Some("My Doc".into()), - doc_type: Some("note".into()), - mime: Some("text/markdown".into()), - tags: Some(vec!["rust".into(), "db".into()]), - source: Some("agent".into()), - ..Default::default() - }).unwrap(); + let doc = s + .doc_upsert_with_opts( + "p", + "/meta.md", + "content", + DocUpsertOpts { + title: Some("My Doc".into()), + doc_type: Some("note".into()), + mime: Some("text/markdown".into()), + tags: Some(vec!["rust".into(), "db".into()]), + source: Some("agent".into()), + ..Default::default() + }, + ) + .unwrap(); assert_eq!(doc.title, "My Doc"); assert_eq!(doc.doc_type, "note"); assert_eq!(doc.mime, "text/markdown"); diff --git a/crates/agentflare-store/src/embedding_pipeline/mod.rs b/crates/agentflare-store/src/embedding_pipeline/mod.rs index 04ebee53..977527d5 100644 --- a/crates/agentflare-store/src/embedding_pipeline/mod.rs +++ b/crates/agentflare-store/src/embedding_pipeline/mod.rs @@ -1,12 +1,13 @@ -pub mod model_registry; pub mod download; -pub mod tokenizer; +pub mod model_registry; pub mod pooling; +pub mod tokenizer; use std::path::{Path, PathBuf}; +use crate::embed; use model_registry::{EmbeddingModel, ModelConfig, VocabSource}; -use tokenizer::{TokenizedInput, WordPieceTokenizer, HfTokenizerWrapper}; +use tokenizer::{HfTokenizerWrapper, TokenizedInput, WordPieceTokenizer}; pub struct EmbeddingEngine { tokenizer: TokenizerKind, @@ -45,11 +46,9 @@ impl EmbeddingEngine { let tokenizer = load_tokenizer(&model_dir, &config)?; let model_path = model_dir.join("model.onnx"); - let session = ort::session::Session::builder() + let session = ort::session::Session::builder() .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))? - .with_intra_threads( - std::thread::available_parallelism().map_or(4, |n| n.get().max(1)), - ) + .with_intra_threads(std::thread::available_parallelism().map_or(4, |n| n.get().max(1))) .map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))? .with_optimization_level(ort::session::builder::GraphOptimizationLevel::All) .map_err(|e| anyhow::anyhow!("ORT optimization: {e}"))? @@ -123,9 +122,14 @@ impl EmbeddingEngine { let input = tokenize(&self.tokenizer, input_text, self.max_seq_len); let mut hidden = self.run_inference(&input)?; - hidden = pooling::mean_pool(&hidden, &input.attention_mask, input.input_ids.len(), self.dimensions); + hidden = pooling::mean_pool( + &hidden, + &input.attention_mask, + input.input_ids.len(), + self.dimensions, + ); - pooling::normalize_l2(&mut hidden); + embed::normalize(&mut hidden); Ok(hidden) } @@ -148,10 +152,14 @@ impl EmbeddingEngine { ); hidden = pooled; - pooling::normalize_l2(&mut hidden); + embed::normalize(&mut hidden); Ok(hidden) } + pub fn model_id(&self) -> &EmbeddingModel { + &self.model_id + } + pub fn dimensions(&self) -> usize { self.dimensions } @@ -164,11 +172,10 @@ impl EmbeddingEngine { if let Ok(dir) = std::env::var("AGENTFLARE_MODELS_DIR") { return PathBuf::from(dir); } - let base = dirs::cache_dir() + dirs::cache_dir() .unwrap_or_else(|| PathBuf::from(".")) .join("agentflare") - .join("models"); - base + .join("models") } fn run_inference(&self, input: &TokenizedInput) -> anyhow::Result> { @@ -274,7 +281,9 @@ fn detect_dimensions( }; let (shape, _) = outputs[output_name].try_extract_tensor::()?; - shape.last().copied().map(|s| s as usize).ok_or_else(|| { - anyhow::anyhow!("could not detect embedding dimensions from model output") - }) + shape + .last() + .copied() + .map(|s| s as usize) + .ok_or_else(|| anyhow::anyhow!("could not detect embedding dimensions from model output")) } diff --git a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs index d9b57612..57073085 100644 --- a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs +++ b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs @@ -26,12 +26,18 @@ impl CustomModelSpec { }; let mut parts = repo.split('/'); let (owner, name) = (parts.next()?, parts.next()?); - if parts.next().is_some() || owner.is_empty() || name.is_empty() + if parts.next().is_some() + || owner.is_empty() + || name.is_empty() || repo.chars().any(char::is_whitespace) { return None; } - Some(Self { repo: repo.to_string(), revision, dimensions: None }) + Some(Self { + repo: repo.to_string(), + revision, + dimensions: None, + }) } fn storage_slug(&self) -> String { diff --git a/crates/agentflare-store/src/embedding_pipeline/pooling.rs b/crates/agentflare-store/src/embedding_pipeline/pooling.rs index a7bf46e5..4db7238f 100644 --- a/crates/agentflare-store/src/embedding_pipeline/pooling.rs +++ b/crates/agentflare-store/src/embedding_pipeline/pooling.rs @@ -10,9 +10,9 @@ pub fn mean_pool( for pos in 0..seq_len { if attention_mask.get(pos).copied().unwrap_or(0) > 0 { let offset = pos * dim; - for d in 0..dim { + for (d, s) in sum.iter_mut().enumerate() { if let Some(&val) = hidden_states.get(offset + d) { - sum[d] += val; + *s += val; } } count += 1.0; @@ -26,16 +26,3 @@ pub fn mean_pool( } sum } - -pub fn normalize_l2(vec: &mut [f32]) { - let norm: f32 = vec.iter().map(|x| x * x).sum::().sqrt(); - if norm > f32::EPSILON { - for x in vec.iter_mut() { - *x /= norm; - } - } -} - -pub fn l2_norm(vec: &[f32]) -> f32 { - vec.iter().map(|x| x * x).sum::().sqrt() -} diff --git a/crates/agentflare-store/src/kv.rs b/crates/agentflare-store/src/kv.rs index bb176ccd..c2b08899 100644 --- a/crates/agentflare-store/src/kv.rs +++ b/crates/agentflare-store/src/kv.rs @@ -27,29 +27,29 @@ impl Store { pub fn kv_get(&self, key: &str) -> rusqlite::Result> { let conn = self.conn(); conn.query_row( - "SELECT key, value, created_at, updated_at FROM store_kv WHERE key = ?1", - params![key], - |row| { - Ok(KvEntry { - key: row.get(0)?, - value: row.get(1)?, - created_at: row.get(2)?, - updated_at: row.get(3)?, - }) - }, - ) - .optional() + "SELECT key, value, created_at, updated_at FROM store_kv WHERE key = ?1", + params![key], + |row| { + Ok(KvEntry { + key: row.get(0)?, + value: row.get(1)?, + created_at: row.get(2)?, + updated_at: row.get(3)?, + }) + }, + ) + .optional() } pub fn kv_exists(&self, key: &str) -> rusqlite::Result { let conn = self.conn(); conn.query_row( - "SELECT 1 FROM store_kv WHERE key = ?1", - params![key], - |_| Ok(()), - ) - .optional() - .map(|o| o.is_some()) + "SELECT 1 FROM store_kv WHERE key = ?1", + params![key], + |_| Ok(()), + ) + .optional() + .map(|o| o.is_some()) } pub fn kv_delete(&self, key: &str) -> rusqlite::Result { From 73c1c92be3b650c6dcca99df809fa1be58a17eea Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 23:58:04 +0530 Subject: [PATCH 07/16] fix(agentflare-store): dispatch tokenizer by model.type, implement BPE decoding --- .../src/embedding_pipeline/tokenizer.rs | 375 ++++++++++++++++-- 1 file changed, 334 insertions(+), 41 deletions(-) diff --git a/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs index ea094382..1cdc4b50 100644 --- a/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs +++ b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs @@ -31,16 +31,27 @@ impl WordPieceTokenizer { .map(|(i, line)| (line.to_string(), i as i32)) .collect(); - let cls_id = *vocab.get("[CLS]") + let cls_id = *vocab + .get("[CLS]") .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [CLS] token"))?; - let sep_id = *vocab.get("[SEP]") + let sep_id = *vocab + .get("[SEP]") .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [SEP] token"))?; - let pad_id = *vocab.get("[PAD]") + let pad_id = *vocab + .get("[PAD]") .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [PAD] token"))?; - let unk_id = *vocab.get("[UNK]") + let unk_id = *vocab + .get("[UNK]") .ok_or_else(|| anyhow::anyhow!("Vocabulary missing [UNK] token"))?; - Ok(Self { vocab, cls_id, sep_id, pad_id, unk_id, max_word_chars: 200 }) + Ok(Self { + vocab, + cls_id, + sep_id, + pad_id, + unk_id, + max_word_chars: 200, + }) } pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { @@ -48,10 +59,14 @@ impl WordPieceTokenizer { let mut ids = vec![self.cls_id]; for word in &words { - if ids.len() >= max_len - 1 { break; } + if ids.len() >= max_len - 1 { + break; + } let subword_ids = self.wordpiece_encode(word); for id in subword_ids { - if ids.len() >= max_len - 1 { break; } + if ids.len() >= max_len - 1 { + break; + } ids.push(id); } } @@ -101,16 +116,28 @@ impl WordPieceTokenizer { let chars: Vec = word.chars().collect(); for (i, &ch) in chars.iter().enumerate() { if ch == '_' || ch == '-' { - if !current.is_empty() { parts.push(current.clone()); current.clear(); } + if !current.is_empty() { + parts.push(current.clone()); + current.clear(); + } } else if i > 0 && ch.is_ascii_uppercase() && chars[i - 1].is_ascii_lowercase() { - if !current.is_empty() { parts.push(current.clone()); current.clear(); } + if !current.is_empty() { + parts.push(current.clone()); + current.clear(); + } current.push(ch); } else { current.push(ch); } } - if !current.is_empty() { parts.push(current); } - if parts.is_empty() { vec![word.to_string()] } else { parts } + if !current.is_empty() { + parts.push(current); + } + if parts.is_empty() { + vec![word.to_string()] + } else { + parts + } } fn wordpiece_encode(&self, word: &str) -> Vec { @@ -125,7 +152,11 @@ impl WordPieceTokenizer { let mut matched = false; while start < end { let substr: String = chars[start..end].iter().collect(); - let candidate = if start > 0 { format!("##{substr}") } else { substr }; + let candidate = if start > 0 { + format!("##{substr}") + } else { + substr + }; if let Some(&id) = self.vocab.get(&candidate) { tokens.push(id); matched = true; @@ -143,8 +174,144 @@ impl WordPieceTokenizer { } } +pub struct BpeTokenizer { + vocab: HashMap, + ranks: HashMap<(String, String), usize>, + unk_id: i32, + lowercase: bool, +} + +impl BpeTokenizer { + fn from_json(model: &serde_json::Value) -> anyhow::Result { + let vocab_obj = model + .get("vocab") + .and_then(|v| v.as_object()) + .ok_or_else(|| anyhow::anyhow!("tokenizer.json BPE model missing vocab object"))?; + + let mut vocab = HashMap::new(); + for (token, id) in vocab_obj { + if let Some(id) = id.as_i64() { + vocab.insert(token.clone(), id as i32); + } + } + + let unk_id = *vocab + .get("") + .or_else(|| vocab.get("")) + .or_else(|| vocab.get("")) + .ok_or_else(|| anyhow::anyhow!("BPE vocab missing token"))?; + + let merges = model + .get("merges") + .and_then(|m| m.as_array()) + .ok_or_else(|| anyhow::anyhow!("tokenizer.json BPE model missing merges"))?; + + let mut ranks = HashMap::new(); + for (i, m) in merges.iter().enumerate() { + let pair = match m { + serde_json::Value::String(s) => { + let mut parts = s.splitn(2, ' '); + let left = parts.next().unwrap_or("").to_string(); + let right = parts.next().unwrap_or("").to_string(); + (left, right) + } + serde_json::Value::Array(arr) if arr.len() == 2 => { + let left = arr[0].as_str().unwrap_or("").to_string(); + let right = arr[1].as_str().unwrap_or("").to_string(); + (left, right) + } + _ => anyhow::bail!("BPE merge entry malformed: {m}"), + }; + ranks.insert(pair, i); + } + + let lowercase = model + .get("lowercase") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + Ok(Self { + vocab, + ranks, + unk_id, + lowercase, + }) + } + + pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { + let lowered: String; + let text = if self.lowercase { + lowered = text.to_lowercase(); + lowered.as_str() + } else { + text + }; + + let mut ids = Vec::new(); + for word in text.split_whitespace() { + self.bpe_encode_word(word, &mut ids, max_len); + if ids.len() >= max_len { + break; + } + } + + let len = ids.len(); + TokenizedInput { + input_ids: ids, + attention_mask: vec![1; len], + token_type_ids: vec![0; len], + } + } + + fn bpe_encode_word(&self, word: &str, ids: &mut Vec, max_len: usize) { + if word.is_empty() { + return; + } + let mut symbols: Vec = word.chars().map(|c| c.to_string()).collect(); + + loop { + if symbols.len() < 2 { + break; + } + let mut best_rank: Option = None; + let mut best_idx: Option = None; + for i in 0..symbols.len() - 1 { + if let Some(&r) = self + .ranks + .get(&(symbols[i].clone(), symbols[i + 1].clone())) + && best_rank.is_none_or(|br| r < br) + { + best_rank = Some(r); + best_idx = Some(i); + } + } + match best_idx { + Some(i) => { + let merged = symbols[i].clone() + &symbols[i + 1]; + symbols.remove(i + 1); + symbols[i] = merged; + } + None => break, + } + } + + for sym in symbols { + if ids.len() >= max_len { + break; + } + let id = self.vocab.get(&sym).copied().unwrap_or(self.unk_id); + ids.push(id); + } + } +} + +pub enum HfTokenizerInner { + WordPiece(WordPieceTokenizer), + Bpe(BpeTokenizer), +} + pub struct HfTokenizerWrapper { - inner: WordPieceTokenizer, + inner: HfTokenizerInner, } impl HfTokenizerWrapper { @@ -155,50 +322,176 @@ impl HfTokenizerWrapper { fn from_json(json_str: &str) -> anyhow::Result { let parsed: serde_json::Value = serde_json::from_str(json_str)?; - let vocab_obj = parsed + let model = parsed .get("model") - .and_then(|m| m.get("vocab")) - .and_then(|v| v.as_object()) - .ok_or_else(|| anyhow::anyhow!("tokenizer.json missing model.vocab object"))?; + .ok_or_else(|| anyhow::anyhow!("tokenizer.json missing model"))?; + let model_type = model + .get("type") + .and_then(|t| t.as_str()) + .ok_or_else(|| anyhow::anyhow!("tokenizer.json model missing type"))?; - let mut vocab_lines: Vec<(String, i32)> = vocab_obj - .iter() - .filter_map(|(token, id)| id.as_i64().map(|id| (token.clone(), id as i32))) - .collect(); - vocab_lines.sort_by_key(|(_, id)| *id); + let inner = match model_type { + "WordPiece" => { + let vocab_obj = model + .get("vocab") + .and_then(|v| v.as_object()) + .ok_or_else(|| anyhow::anyhow!("tokenizer.json missing model.vocab object"))?; - for (token, _) in &mut vocab_lines { - let mapped = match token.as_str() { - "" => "[CLS]", "" => "[SEP]", - "" => "[PAD]", "" => "[UNK]", - "" => "[MASK]", _ => continue, - }; - *token = mapped.to_string(); - } + let mut vocab_lines: Vec<(String, i32)> = vocab_obj + .iter() + .filter_map(|(token, id)| id.as_i64().map(|id| (token.clone(), id as i32))) + .collect(); + vocab_lines.sort_by_key(|(_, id)| *id); + + for (token, _) in &mut vocab_lines { + let mapped = match token.as_str() { + "" => "[CLS]", + "" => "[SEP]", + "" => "[PAD]", + "" => "[UNK]", + "" => "[MASK]", + _ => continue, + }; + *token = mapped.to_string(); + } + + let vocab_str: String = vocab_lines + .into_iter() + .map(|(token, _)| token) + .collect::>() + .join("\n"); - let vocab_str: String = vocab_lines.into_iter() - .map(|(token, _)| token) - .collect::>() - .join("\n"); + let wp = WordPieceTokenizer::from_vocab_str(&vocab_str)?; + HfTokenizerInner::WordPiece(wp) + } + "BPE" => { + let bpe = BpeTokenizer::from_json(model)?; + HfTokenizerInner::Bpe(bpe) + } + other => anyhow::bail!( + "Unsupported tokenizer model.type: {other}. Only WordPiece and BPE are supported." + ), + }; - let inner = WordPieceTokenizer::from_vocab_str(&vocab_str)?; Ok(Self { inner }) } pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { - self.inner.encode(text, max_len) + match &self.inner { + HfTokenizerInner::WordPiece(wp) => wp.encode(text, max_len), + HfTokenizerInner::Bpe(bpe) => bpe.encode(text, max_len), + } } } fn is_bert_punctuation(ch: char) -> bool { if ch.is_ascii() { - matches!(ch, - '!' | '"' | '#' | '$' | '%' | '&' | '\'' | '(' | ')' - | '*' | '+' | ',' | '-' | '.' | '/' | ':' | ';' - | '<' | '=' | '>' | '?' | '@' | '[' | '\\' | ']' - | '^' | '_' | '`' | '{' | '|' | '}' | '~' + matches!( + ch, + '!' | '"' + | '#' + | '$' + | '%' + | '&' + | '\'' + | '(' + | ')' + | '*' + | '+' + | ',' + | '-' + | '.' + | '/' + | ':' + | ';' + | '<' + | '=' + | '>' + | '?' + | '@' + | '[' + | '\\' + | ']' + | '^' + | '_' + | '`' + | '{' + | '|' + | '}' + | '~' ) } else { ch.is_ascii_punctuation() } } + +#[cfg(test)] +mod tests { + use super::*; + + const BPE_JSON: &str = r#"{ + "model": { + "type": "BPE", + "vocab": { + "": 0, + "h": 1, "e": 2, "l": 3, "o": 4, + "he": 5, "hel": 6, "hell": 7, "hello": 8 + }, + "merges": ["h e", "he l", "hel l", "hell o"], + "unk_token": "", + "lowercase": true + } + }"#; + + const WORDPICE_JSON: &str = r###"{ + "model": { + "type": "WordPiece", + "vocab": { + "[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, + "hello": 4, "##llo": 5, "world": 6 + } + } + }"###; + + #[test] + fn bpe_applies_merge_rules() { + let tok = HfTokenizerWrapper::from_json(BPE_JSON).unwrap(); + let out = tok.encode("hello", 32); + assert_eq!(out.input_ids, vec![8]); + } + + #[test] + fn bpe_tokenizes_multiple_words_and_unk() { + let tok = HfTokenizerWrapper::from_json(BPE_JSON).unwrap(); + let out = tok.encode("hello hello", 32); + assert_eq!(out.input_ids, vec![8, 8]); + + let out = tok.encode("helloz", 32); + assert_eq!(out.input_ids, vec![8, 0]); + } + + #[test] + fn wordpiece_path_still_works() { + let tok = HfTokenizerWrapper::from_json(WORDPICE_JSON).unwrap(); + let out = tok.encode("hello", 32); + // [CLS]=2, hello=4, [SEP]=3 + assert_eq!(out.input_ids, vec![2, 4, 3]); + } + + #[test] + fn unknown_model_type_fails_loudly() { + let json = r#"{ "model": { "type": "Unigram", "vocab": {} } }"#; + let err = HfTokenizerWrapper::from_json(json); + let msg = match err { + Ok(_) => panic!("expected Err for unsupported model.type"), + Err(e) => e.to_string(), + }; + assert!(msg.contains("Unsupported tokenizer")); + } + + #[test] + fn missing_model_type_fails() { + let json = r#"{ "model": { "vocab": {} } }"#; + assert!(HfTokenizerWrapper::from_json(json).is_err()); + } +} From 4f592d120f5b4c9fcdc283f636aadade98ade6fc Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Fri, 17 Jul 2026 23:55:57 +0530 Subject: [PATCH 08/16] fix(agentflare-store): re-verify size+SHA-256 of already-present model files --- .../src/embedding_pipeline/download.rs | 155 +++++++++++++++++- .../src/embedding_pipeline/model_registry.rs | 9 + 2 files changed, 161 insertions(+), 3 deletions(-) diff --git a/crates/agentflare-store/src/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs index 58b7aa3f..b950d7b8 100644 --- a/crates/agentflare-store/src/embedding_pipeline/download.rs +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -16,14 +16,35 @@ struct DownloadFile { pub fn ensure_model(model_dir: &Path, config: &ModelConfig) -> anyhow::Result { let files = download_files(config); - let all_present = files.iter().all(|f| model_dir.join(&f.local_name).exists()); - if all_present { + // Re-verify already-present files (size + SHA-256 against the pinned hash in + // model.lock.json) before trusting them. A present-but-corrupt file (partial + // write, disk corruption, tampering) must not silently bypass the checksum, + // so a size/hash mismatch is treated as missing and re-downloaded below. + let lock = read_lockfile(model_dir); + let mut any_corrupt = false; + for file in &files { + let local_path = model_dir.join(&file.local_name); + if !local_path.exists() { + continue; + } + if !file_passes_verification(&local_path, file, &lock)? { + tracing::warn!( + "Embedding model file {} present but failed size/SHA-256 verification; re-downloading", + file.local_name + ); + let _ = std::fs::remove_file(&local_path); + any_corrupt = true; + } + } + + let all_present = files.iter().all(|f| model_dir.join(&f.local_name).exists()); + if all_present && !any_corrupt { return Ok(model_dir.to_path_buf()); } tracing::info!( - "Embedding model '{}' not found, downloading to {}", + "Embedding model '{}' not found or invalid, downloading to {}", config.name, model_dir.display() ); @@ -130,6 +151,25 @@ fn download_file( Ok(()) } +/// Re-verify a present model file: it must meet the minimum size and, when a +/// pinned SHA-256 exists in the lockfile, match that hash. Returns `true` when +/// the file can be trusted as-is (so `ensure_model` need not re-download it). +fn file_passes_verification( + local_path: &Path, + file: &DownloadFile, + lock: &BTreeMap, +) -> anyhow::Result { + let meta = std::fs::metadata(local_path)?; + if meta.len() < file.min_bytes { + return Ok(false); + } + let actual = sha256_file(local_path)?; + match lock.get(&file.local_name) { + Some(pinned) if pinned != &actual => Ok(false), + _ => Ok(true), + } +} + fn sha256_file(path: &Path) -> anyhow::Result { use sha2::{Digest, Sha256}; let mut file = std::fs::File::open(path)?; @@ -175,3 +215,112 @@ pub fn clean_model(model_dir: &Path) -> anyhow::Result<()> { } Ok(()) } + +#[cfg(all(test, feature = "embeddings"))] +mod tests { + use crate::embedding_pipeline::model_registry::{EmbeddingModel, ModelConfig}; + use super::*; + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + const MODEL_BODY: &[u8] = b"MODEL-FILE-CONTENTS-0123456789-abcdefghij"; + const VOCAB_BODY: &[u8] = b"vocab line one\nvocab line two\nvocab line three\n"; + + /// Serves two fixed files from a local HTTP server so the regression test + /// needs no network and produces deterministic SHA-256 hashes. + fn spawn_server() -> (String, Arc) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + let hits = Arc::new(AtomicUsize::new(0)); + let hits2 = hits.clone(); + std::thread::spawn(move || { + for stream in listener.incoming() { + let mut stream = match stream { + Ok(s) => s, + Err(_) => continue, + }; + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf); + let body: &[u8] = if buf.starts_with(b"GET /onnx/model.onnx") { + MODEL_BODY + } else { + VOCAB_BODY + }; + hits2.fetch_add(1, Ordering::SeqCst); + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + let _ = stream.write_all(header.as_bytes()); + let _ = stream.write_all(body); + let _ = stream.flush(); + } + }); + (format!("http://127.0.0.1:{port}"), hits) + } + + fn test_config(base: &str) -> ModelConfig { + let mut cfg = EmbeddingModel::AllMiniLmL6V2.config(); + cfg.base_url_override = Some(base.to_string()); + cfg.model_min_bytes = 1; + cfg.vocab_min_bytes = 1; + cfg + } + + fn sha256(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut h = Sha256::new(); + h.update(bytes); + format!("{:x}", h.finalize()) + } + + #[test] + fn present_file_reverified_on_hash_mismatch() { + let (base, hits) = spawn_server(); + let dir = tempfile::tempdir().unwrap(); + let model_dir = dir.path().join("model"); + let config = test_config(&base); + + // First call downloads and pins both files. + ensure_model(&model_dir, &config).unwrap(); + assert_eq!(hits.load(Ordering::SeqCst), 2); + assert_eq!(std::fs::read(model_dir.join("model.onnx")).unwrap(), MODEL_BODY); + assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + + // Tamper with the vocab file so its SHA-256 no longer matches the pin. + std::fs::write(model_dir.join("vocab.txt"), b"CORRUPTED-VOCAB-DATA-NOT-REAL").unwrap(); + + // Second call must detect the mismatch and re-download the file. + ensure_model(&model_dir, &config).unwrap(); + + // Re-downloaded file now matches the genuine server content. + assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + + // The lockfile still records the genuine pinned hashes. + let lock: std::collections::BTreeMap = serde_json::from_str( + &std::fs::read_to_string(model_dir.join("model.lock.json")).unwrap(), + ) + .unwrap(); + assert_eq!(lock.get("vocab.txt").unwrap(), &sha256(VOCAB_BODY)); + assert_eq!(lock.get("model.onnx").unwrap(), &sha256(MODEL_BODY)); + assert!(hits.load(Ordering::SeqCst) >= 3); + } + + #[test] + fn present_file_reverified_on_size_mismatch() { + let (base, _hits) = spawn_server(); + let dir = tempfile::tempdir().unwrap(); + let model_dir = dir.path().join("model"); + let config = test_config(&base); + + ensure_model(&model_dir, &config).unwrap(); + + // Truncate the vocab file below the minimum size. + std::fs::write(model_dir.join("vocab.txt"), b"").unwrap(); + + ensure_model(&model_dir, &config).unwrap(); + assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + } +} diff --git a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs index 57073085..319ac5e3 100644 --- a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs +++ b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs @@ -82,6 +82,7 @@ impl EmbeddingModel { query_prefix: None, document_prefix: None, needs_token_type_ids: true, + base_url_override: None, }, Self::NomicEmbedV1_5 => ModelConfig { model: self.clone(), @@ -97,6 +98,7 @@ impl EmbeddingModel { query_prefix: Some("search_query: ".into()), document_prefix: Some("search_document: ".into()), needs_token_type_ids: false, + base_url_override: None, }, Self::Custom(spec) => ModelConfig { model: self.clone(), @@ -115,6 +117,7 @@ impl EmbeddingModel { query_prefix: None, document_prefix: None, needs_token_type_ids: false, + base_url_override: None, }, } } @@ -172,6 +175,9 @@ pub struct ModelConfig { pub revision: Option, pub onnx_path: String, pub vocab_file: VocabSource, + /// When set, overrides the HuggingFace base URL used to resolve downloads. + /// Primarily for tests that serve model files from a local HTTP server. + pub base_url_override: Option, pub dimensions: usize, pub max_seq_len: usize, pub model_min_bytes: u64, @@ -183,6 +189,9 @@ pub struct ModelConfig { impl ModelConfig { fn resolve_base(&self) -> String { + if let Some(base) = &self.base_url_override { + return base.clone(); + } format!( "https://huggingface.co/{}/resolve/{}", self.hf_repo, From 0cf7fa00489b76b9d9a9dbae8e0f3f9b5e2e9489 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 00:04:40 +0530 Subject: [PATCH 09/16] fix(agentflare-store): cargo fmt download.rs after merging #153's fix --- .../src/embedding_pipeline/download.rs | 43 ++++++++++++++----- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/crates/agentflare-store/src/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs index b950d7b8..06e33f0e 100644 --- a/crates/agentflare-store/src/embedding_pipeline/download.rs +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -68,7 +68,8 @@ pub fn ensure_model(model_dir: &Path, config: &ModelConfig) -> anyhow::Result anyhow::Result= {min_bytes})"); + anyhow::bail!( + "Downloaded {local_name} is too small ({total} bytes, expected >= {min_bytes})" + ); } std::fs::rename(&tmp_path, &local_path)?; @@ -218,12 +225,12 @@ pub fn clean_model(model_dir: &Path) -> anyhow::Result<()> { #[cfg(all(test, feature = "embeddings"))] mod tests { - use crate::embedding_pipeline::model_registry::{EmbeddingModel, ModelConfig}; use super::*; + use crate::embedding_pipeline::model_registry::{EmbeddingModel, ModelConfig}; use std::io::{Read, Write}; use std::net::TcpListener; - use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; const MODEL_BODY: &[u8] = b"MODEL-FILE-CONTENTS-0123456789-abcdefghij"; const VOCAB_BODY: &[u8] = b"vocab line one\nvocab line two\nvocab line three\n"; @@ -286,17 +293,30 @@ mod tests { // First call downloads and pins both files. ensure_model(&model_dir, &config).unwrap(); assert_eq!(hits.load(Ordering::SeqCst), 2); - assert_eq!(std::fs::read(model_dir.join("model.onnx")).unwrap(), MODEL_BODY); - assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + assert_eq!( + std::fs::read(model_dir.join("model.onnx")).unwrap(), + MODEL_BODY + ); + assert_eq!( + std::fs::read(model_dir.join("vocab.txt")).unwrap(), + VOCAB_BODY + ); // Tamper with the vocab file so its SHA-256 no longer matches the pin. - std::fs::write(model_dir.join("vocab.txt"), b"CORRUPTED-VOCAB-DATA-NOT-REAL").unwrap(); + std::fs::write( + model_dir.join("vocab.txt"), + b"CORRUPTED-VOCAB-DATA-NOT-REAL", + ) + .unwrap(); // Second call must detect the mismatch and re-download the file. ensure_model(&model_dir, &config).unwrap(); // Re-downloaded file now matches the genuine server content. - assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + assert_eq!( + std::fs::read(model_dir.join("vocab.txt")).unwrap(), + VOCAB_BODY + ); // The lockfile still records the genuine pinned hashes. let lock: std::collections::BTreeMap = serde_json::from_str( @@ -321,6 +341,9 @@ mod tests { std::fs::write(model_dir.join("vocab.txt"), b"").unwrap(); ensure_model(&model_dir, &config).unwrap(); - assert_eq!(std::fs::read(model_dir.join("vocab.txt")).unwrap(), VOCAB_BODY); + assert_eq!( + std::fs::read(model_dir.join("vocab.txt")).unwrap(), + VOCAB_BODY + ); } } From 9ed6d6b91c8e240782669a5593c2c1f6c604c1d4 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 00:20:51 +0530 Subject: [PATCH 10/16] feat(agentflare-store): wire state.json onto the store's kv table (Task 8) src/store.rs::open() previously cached the store behind a OnceLock -- dead code (nothing called it) and would have broken test isolation the moment something did, since AGENTFLARE_HOME_OVERRIDE changes per-test but a cached singleton would keep pointing at whichever home dir opened it first. Open fresh per call instead, matching memory::store::open()'s existing pattern in this codebase. state::load()/save() now read/write through agentflare-store's kv table under the active/version_cache keys, with a one-time import of any legacy state.json via agentflare_store::migrate::migrate_state_json on first load against a store that has neither key yet. Public API unchanged (State, load(), save(), state_path()), so no callers change. --- src/state.rs | 108 +++++++++++++++++++++++++++++++++++++++++++-------- src/store.rs | 43 ++++++++++++-------- 2 files changed, 118 insertions(+), 33 deletions(-) diff --git a/src/state.rs b/src/state.rs index e364356f..1b4fc5aa 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,10 +1,12 @@ -// Single JSON state blob, host-neutral (~/.agentflare/), shared across -// whichever agents this machine has run `agentflare init`/hooks for. +// Host-neutral local state (~/.agentflare/), shared across whichever agents +// this machine has run `agentflare init`/hooks for. Backed by +// `agentflare-store`'s kv table; a legacy `state.json` (the pre-store +// on-disk format) is imported once, in place, the first time this runs +// against a store that has neither key yet. use crate::paths::home; pub use agent_registry::VersionCacheEntry; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::fs; use std::path::PathBuf; #[derive(Serialize, Deserialize, Default)] @@ -28,24 +30,79 @@ pub fn state_path() -> PathBuf { state_dir().join("state.json") } +const ACTIVE_KEY: &str = "active"; +const VERSION_CACHE_KEY: &str = "version_cache"; + +fn default_state() -> State { + State { + active: true, + ..Default::default() + } +} + pub fn load() -> State { - fs::read_to_string(state_path()) + let store = match crate::store::open() { + Ok(store) => store, + Err(e) => { + eprintln!("[agentflare] warning: failed to open store: {e}"); + return default_state(); + } + }; + + // One-time bridge for installs that predate the store: if neither key + // is present yet, but a legacy state.json is on disk, import it. + // migrate_state_json flattens state.json's top-level keys ("active", + // "version_cache") straight into these same kv keys, and records its + // own marker so it never re-runs. A parse failure on a corrupt legacy + // file is swallowed here -- the kv reads below then fall through to + // defaults, same as the old file-based load() did on corrupt JSON. + let has_active = store.kv_get(ACTIVE_KEY).ok().flatten().is_some(); + let has_version_cache = store.kv_get(VERSION_CACHE_KEY).ok().flatten().is_some(); + if !has_active && !has_version_cache { + let legacy_path = state_path(); + if legacy_path.exists() { + let _ = agentflare_store::migrate::migrate_state_json(&store, &legacy_path); + } + } + + let active = store + .kv_get(ACTIVE_KEY) + .ok() + .flatten() + .and_then(|entry| serde_json::from_slice(&entry.value).ok()) + .unwrap_or(true); + let version_cache = store + .kv_get(VERSION_CACHE_KEY) .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or_else(|| State { - active: true, - ..Default::default() - }) + .flatten() + .and_then(|entry| serde_json::from_slice(&entry.value).ok()) + .unwrap_or_default(); + + State { + active, + version_cache, + } } pub fn save(state: &State) { - if let Err(e) = fs::create_dir_all(state_dir()) { - eprintln!("[agentflare] warning: failed to create state dir: {e}"); - return; + let store = match crate::store::open() { + Ok(store) => store, + Err(e) => { + eprintln!("[agentflare] warning: failed to open store: {e}"); + return; + } + }; + match serde_json::to_vec(&state.active) { + Ok(bytes) => { + if let Err(e) = store.kv_set(ACTIVE_KEY, &bytes) { + eprintln!("[agentflare] warning: failed to persist state: {e}"); + } + } + Err(e) => eprintln!("[agentflare] warning: failed to serialize state: {e}"), } - match serde_json::to_string_pretty(state) { - Ok(json) => { - if let Err(e) = fs::write(state_path(), json + "\n") { + match serde_json::to_vec(&state.version_cache) { + Ok(bytes) => { + if let Err(e) = store.kv_set(VERSION_CACHE_KEY, &bytes) { eprintln!("[agentflare] warning: failed to persist state: {e}"); } } @@ -57,6 +114,7 @@ pub fn save(state: &State) { mod tests { use super::*; use crate::paths::test_support::with_temp_home; + use std::fs; #[test] fn load_defaults_to_active_when_no_state_file() { @@ -77,7 +135,7 @@ mod tests { } #[test] - fn load_falls_back_to_default_on_corrupt_file() { + fn load_falls_back_to_default_on_corrupt_legacy_file() { with_temp_home(|| { fs::create_dir_all(state_dir()).unwrap(); fs::write(state_path(), "not json").unwrap(); @@ -124,4 +182,22 @@ mod tests { assert!(s.version_cache.is_empty()); }); } + + #[test] + fn legacy_state_file_is_migrated_exactly_once() { + with_temp_home(|| { + fs::create_dir_all(state_dir()).unwrap(); + fs::write(state_path(), r#"{"active": false}"#).unwrap(); + + let first = load(); + assert!(!first.active); + + // Mutate the legacy file after the first load -- since + // migration already ran (kv keys now present), the second + // load must read from the store, not re-import the file. + fs::write(state_path(), r#"{"active": true}"#).unwrap(); + let second = load(); + assert!(!second.active); + }); + } } diff --git a/src/store.rs b/src/store.rs index d4cbaf41..6ebf7d21 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,24 +1,24 @@ -use agentflare_store::Store; +use agentflare_store::{Error, Store}; use std::path::PathBuf; -use std::sync::OnceLock; - -static STORE: OnceLock> = OnceLock::new(); pub fn store_path() -> PathBuf { - crate::paths::home().join(".agentflare").join("agentflare.db") + crate::paths::home() + .join(".agentflare") + .join("agentflare.db") } -pub fn open() -> &'static Store { - STORE - .get_or_init(|| { - let path = store_path(); - if let Some(parent) = path.parent() { - let _ = std::fs::create_dir_all(parent); - } - Store::open_file(&path) - }) - .as_ref() - .expect("failed to open agentflare store") +/// Opens a fresh connection to the local store on every call -- deliberately +/// not cached behind a `OnceLock`. `crate::paths::home()` respects +/// `AGENTFLARE_HOME_OVERRIDE` (see `paths::test_support::with_temp_home`), +/// which tests flip per-call; a cached singleton would keep pointing at the +/// first test's home dir forever. Mirrors `memory::store::open()`'s same +/// per-call-open pattern for the same reason. +pub fn open() -> Result { + let path = store_path(); + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + Store::open_file(&path) } #[cfg(test)] @@ -29,8 +29,17 @@ mod tests { #[test] fn open_and_close() { with_temp_home(|| { - let store = open(); + let store = open().unwrap(); store.conn().execute_batch("SELECT 1").unwrap(); }); } + + #[test] + fn each_call_sees_the_current_home_override() { + with_temp_home(|| { + let first = store_path(); + open().unwrap(); + assert!(first.exists()); + }); + } } From 6822d0784f54882bb204a17f642b7783c52f19a5 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 00:34:13 +0530 Subject: [PATCH 11/16] fix(agentflare-store): stop colliding with db.rs's agentflare.db, use store.db store_path() picked ~/.agentflare/agentflare.db -- the same file src/db.rs already owns as its single source-of-truth relational store (claims, handoffs, review_findings, gateway_secrets), with its own separate migration list. Two independent migration systems targeting one file is a real hazard, not just an aesthetic collision. agentflare-store is a different kind of storage (blobs, FTS+vector docs, kv) so a separate file is correct -- it just needed a name that doesn't collide. --- src/store.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/store.rs b/src/store.rs index 6ebf7d21..bcc8b427 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,10 +1,14 @@ use agentflare_store::{Error, Store}; use std::path::PathBuf; +/// Deliberately NOT `agentflare.db` -- that file is `src/db.rs`'s "single +/// source-of-truth" relational store (claims, handoffs, review_findings, +/// gateway_secrets, ...), with its own separate migration list. This store +/// is a different kind of thing (blobs, FTS+vector documents, kv) with its +/// own migrations; sharing a file would let two independent migration +/// systems fight over the same schema/version state. pub fn store_path() -> PathBuf { - crate::paths::home() - .join(".agentflare") - .join("agentflare.db") + crate::paths::home().join(".agentflare").join("store.db") } /// Opens a fresh connection to the local store on every call -- deliberately From 94f37161be81bed568ef0db60ae0df887a9c8f7e Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 12:28:05 +0530 Subject: [PATCH 12/16] fix(agentflare-store): make blob/doc/migrate writes atomic and race-free - blob_store/blob_unref: wrap the exists-check + insert/decrement + cascade in an Immediate transaction so concurrent connections can't race on the same hash; compensate for orphaned disk writes on failure. - doc_upsert_with_opts: open the transaction before reading the current version so concurrent upserts can't both read version N and both write N+1; also makes new-document insert + FTS sync atomic. - doc_hard_delete: delete history/vector/FTS rows before the parent row, all in one transaction. - migrate_state_json: one transaction for every key plus the completion marker, so a partial failure can't desync the marker from what migrated. --- crates/agentflare-store/src/blobs.rs | 72 ++++++++++++++++++------ crates/agentflare-store/src/documents.rs | 65 +++++++++++++-------- crates/agentflare-store/src/migrate.rs | 25 ++++++-- 3 files changed, 118 insertions(+), 44 deletions(-) diff --git a/crates/agentflare-store/src/blobs.rs b/crates/agentflare-store/src/blobs.rs index c89aed51..e5287d19 100644 --- a/crates/agentflare-store/src/blobs.rs +++ b/crates/agentflare-store/src/blobs.rs @@ -17,9 +17,15 @@ fn blob_disk_path(root: &Path, hash: &str) -> PathBuf { root.join("blobs").join(&hash[..2]).join(hash) } -fn read_disk_blob(root: &Path, hash: &str) -> Option> { +/// `Ok(None)` means the file genuinely isn't there; other I/O errors (permissions, +/// disk failures) are propagated instead of being folded into "not found". +fn read_disk_blob(root: &Path, hash: &str) -> std::io::Result>> { let path = blob_disk_path(root, hash); - std::fs::read(&path).ok() + match std::fs::read(&path) { + Ok(data) => Ok(Some(data)), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(e) => Err(e), + } } fn write_disk_blob(root: &Path, hash: &str, data: &[u8]) -> Result<(), std::io::Error> { @@ -47,8 +53,16 @@ impl Store { let conn = self.conn(); let hash = blake3::hash(data).to_hex().to_string(); let now = db_kit::ids::now(); + let is_memory = self.is_memory(); - let exists = conn + // Immediate takes the write lock up front, so the exists-check and + // the insert-or-bump below are atomic across connections — without + // this, two concurrent stores of the same new content can both see + // "not found" and then race on the INSERT. + let tx = + rusqlite::Transaction::new_unchecked(&conn, rusqlite::TransactionBehavior::Immediate)?; + + let exists = tx .query_row( "SELECT 1 FROM store_blobs WHERE hash = ?1", params![hash], @@ -58,31 +72,40 @@ impl Store { .is_some(); if exists { - conn.execute( + tx.execute( "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", params![hash], )?; + tx.commit()?; return Ok(hash); } - let is_memory = self.is_memory(); if !is_memory { + // Written outside the SQL transaction (files aren't part of it); + // if the metadata insert below fails, remove it again so a + // failed store doesn't leak an orphaned file with no DB row. if let Err(e) = write_disk_blob(&self.root, &hash, data) { return Err(rusqlite::Error::ToSqlConversionFailure(Box::new(e))); } } else { for (i, chunk) in data.chunks(CHUNK_SIZE).enumerate() { - conn.execute( + tx.execute( "INSERT INTO store_blob_chunks (hash, chunk_index, data) VALUES (?1, ?2, ?3)", params![hash, i as i64, chunk], )?; } } - conn.execute( + if let Err(e) = tx.execute( "INSERT INTO store_blobs (hash, size, ref_count, created_at) VALUES (?1, ?2, 1, ?3)", params![hash, data.len() as i64, now], - )?; + ) { + if !is_memory { + delete_disk_blob(&self.root, &hash); + } + return Err(e); + } + tx.commit()?; Ok(hash) } @@ -109,7 +132,8 @@ impl Store { }; if !self.is_memory() { - return Ok(read_disk_blob(&self.root, hash)); + return read_disk_blob(&self.root, hash) + .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e))); } let conn = self.conn(); @@ -136,12 +160,22 @@ impl Store { pub fn blob_unref(&self, hash: &str) -> rusqlite::Result { let conn = self.conn(); - let n = conn.execute( + let is_memory = self.is_memory(); + + // Immediate takes the write lock up front, so the decrement and the + // ref_count<=0 cascade-delete below are atomic across connections — + // without this, two concurrent unrefs can both observe ref_count<=0 + // and both attempt the cascade. + let tx = + rusqlite::Transaction::new_unchecked(&conn, rusqlite::TransactionBehavior::Immediate)?; + + let n = tx.execute( "UPDATE store_blobs SET ref_count = ref_count - 1 WHERE hash = ?1 AND ref_count > 0", params![hash], )?; + let mut removed = false; if n > 0 { - let removed = conn + removed = tx .query_row( "SELECT ref_count <= 0 FROM store_blobs WHERE hash = ?1", params![hash], @@ -150,17 +184,23 @@ impl Store { .optional()? .unwrap_or(false); if removed { - conn.execute("DELETE FROM store_blobs WHERE hash = ?1", params![hash])?; - if !self.is_memory() { - delete_disk_blob(&self.root, hash); - } else { - conn.execute( + if is_memory { + tx.execute( "DELETE FROM store_blob_chunks WHERE hash = ?1", params![hash], )?; } + tx.execute("DELETE FROM store_blobs WHERE hash = ?1", params![hash])?; } } + tx.commit()?; + + // Disk cleanup runs after the metadata commit: the row is the + // source of truth and is already gone, so a crash here just leaks + // a file instead of leaving a dangling row with no data. + if removed && !is_memory { + delete_disk_blob(&self.root, hash); + } Ok(n > 0) } } diff --git a/crates/agentflare-store/src/documents.rs b/crates/agentflare-store/src/documents.rs index 76188d38..ea55e0f1 100644 --- a/crates/agentflare-store/src/documents.rs +++ b/crates/agentflare-store/src/documents.rs @@ -111,7 +111,14 @@ impl Store { let conn = self.conn(); let now = db_kit::ids::now(); - let existing = conn + // BEGIN IMMEDIATE takes SQLite's write lock up front, so the version + // read below is serialized against other connections instead of + // racing them (two connections could otherwise both read version N + // and both compute N+1). + let tx = + rusqlite::Transaction::new_unchecked(&conn, rusqlite::TransactionBehavior::Immediate)?; + + let existing = tx .query_row( "SELECT id, rowid, content, version, blob_hash, mime FROM store_documents WHERE project_id = ?1 AND path = ?2", @@ -135,8 +142,6 @@ impl Store { let new_version = old_version + 1; let history_id = db_kit::ids::new_id(); - let tx = conn.unchecked_transaction()?; - // Snapshot current version to history tx.execute( "INSERT INTO store_doc_history (id, doc_id, version, content, blob_hash, mime, title, created_at) @@ -210,7 +215,9 @@ impl Store { let tags_json = serde_json::to_string(&tags_val).unwrap_or_else(|_| "[]".to_string()); let source = opts.source.unwrap_or_default(); - conn.execute( + // Insert + FTS sync share this transaction so a failure between + // the two can't leave a document without its search index row. + tx.execute( "INSERT INTO store_documents (id, project_id, path, content, title, doc_type, blob_hash, mime, tags, session_id, source, version, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, 1, ?12, ?12)", @@ -219,8 +226,9 @@ impl Store { mime, tags_json, opts.session_id, source, now ], )?; - let rowid = conn.last_insert_rowid(); - Self::doc_sync_fts(&conn, rowid, content)?; + let rowid = tx.last_insert_rowid(); + Self::doc_sync_fts(&tx, rowid, content)?; + tx.commit()?; Ok(Document { id, project_id: project_id.to_string(), @@ -280,23 +288,32 @@ impl Store { pub fn doc_hard_delete(&self, id: &str) -> rusqlite::Result { let conn = self.conn(); - if let Some(rowid) = conn + let tx = + rusqlite::Transaction::new_unchecked(&conn, rusqlite::TransactionBehavior::Immediate)?; + let Some(rowid) = tx .query_row( "SELECT rowid FROM store_documents WHERE id = ?1", params![id], |row| row.get::<_, i64>(0), ) .optional()? - { - conn.execute("DELETE FROM store_documents WHERE id = ?1", params![id])?; - conn.execute( - "DELETE FROM store_docs_fts WHERE rowid = ?1", - params![rowid], - )?; - Ok(true) - } else { - Ok(false) - } + else { + return Ok(false); + }; + + // Delete dependents before the parent row, all in one transaction. + tx.execute( + "DELETE FROM store_doc_history WHERE doc_id = ?1", + params![id], + )?; + tx.execute("DELETE FROM store_docs_vec WHERE doc_id = ?1", params![id])?; + tx.execute( + "DELETE FROM store_docs_fts WHERE rowid = ?1", + params![rowid], + )?; + tx.execute("DELETE FROM store_documents WHERE id = ?1", params![id])?; + tx.commit()?; + Ok(true) } pub fn doc_history(&self, doc_id: &str) -> rusqlite::Result> { @@ -440,7 +457,7 @@ impl Store { .chunks_exact(4) .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect(); - let sim = crate::embed::cosine_similarity(query_vec, &doc_vec) as f64; + let sim = crate::embed::cosine_similarity(query_vec, &doc_vec)? as f64; Some(( sim, DocMatch { @@ -622,15 +639,19 @@ mod tests { let d2 = s.doc_upsert("p", "/dog.md", "about dogs").unwrap(); let d3 = s.doc_upsert("p", "/car.md", "about cars").unwrap(); - test_embed(&s, &d1.id, 4, 1.0); - test_embed(&s, &d2.id, 4, 0.8); - test_embed(&s, &d3.id, 4, 0.0); + // Directionally distinct so cosine similarity actually differs — + // uniform-value vectors like [1,1,1,1] vs [0.8,0.8,0.8,0.8] are + // collinear and score identically regardless of magnitude. + s.doc_set_embedding(&d1.id, &[1.0, 0.0, 0.0, 0.0]).unwrap(); + s.doc_set_embedding(&d2.id, &[1.0, 1.0, 0.0, 0.0]).unwrap(); + s.doc_set_embedding(&d3.id, &[0.0, 1.0, 0.0, 0.0]).unwrap(); - let query = vec![1.0; 4]; + let query = vec![1.0, 0.0, 0.0, 0.0]; let results = s.doc_vec_search("p", &query, 10).unwrap(); assert_eq!(results.len(), 3); assert_eq!(results[0].id, d1.id); assert_eq!(results[1].id, d2.id); + assert_eq!(results[2].id, d3.id); } #[test] diff --git a/crates/agentflare-store/src/migrate.rs b/crates/agentflare-store/src/migrate.rs index 313d6df2..e2aaa4cc 100644 --- a/crates/agentflare-store/src/migrate.rs +++ b/crates/agentflare-store/src/migrate.rs @@ -24,15 +24,28 @@ pub fn migrate_state_json(store: &Store, path: &Path) -> Result = serde_json::from_str(&content)?; + // One transaction for every key plus the completion marker: a partial + // failure rolls back entirely instead of leaving the marker set (or + // absent) out of sync with which keys actually landed. + let conn = store.conn(); + let tx = rusqlite::Transaction::new_unchecked(&conn, rusqlite::TransactionBehavior::Immediate)?; + let now = db_kit::ids::now(); + let upsert = |key: &str, value: &[u8]| -> Result<(), MigrateError> { + tx.execute( + "INSERT INTO store_kv (key, value, created_at, updated_at) + VALUES (?1, ?2, ?3, ?3) + ON CONFLICT(key) DO UPDATE SET value = ?2, updated_at = ?3", + rusqlite::params![key, value, now], + )?; + Ok(()) + }; + for (key, value) in &map { - let serialized = serde_json::to_vec(value)?; - store.kv_set(key, &serialized)?; + upsert(key, &serde_json::to_vec(value)?)?; } + upsert(MIGRATION_MARKER, &serde_json::to_vec(&now)?)?; - let now = db_kit::ids::now(); - let marker = serde_json::to_vec(&now)?; - store.kv_set(MIGRATION_MARKER, &marker)?; - + tx.commit()?; Ok(map.len()) } From 00d18faa8e996e0ccdc4575e11882b096d91236a Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 12:28:18 +0530 Subject: [PATCH 13/16] fix(agentflare-store): surface errors instead of silently swallowing them - cosine_similarity: return None on mismatched vector lengths instead of letting zip() silently drop the excess coordinates. - blob_get: propagate disk read I/O errors (permissions, disk failure) instead of folding them into 'blob not found'. - read_lockfile: distinguish a missing lockfile (fine, first run) from an unreadable or malformed one (error) instead of treating both as 'no pins', which made every present file look unverified and skip the tamper check. - ensure_model: abort if a corrupt model file can't be deleted, instead of proceeding while the corrupt file is still on disk. - resolve_model: error on an invalid AGENTFLARE_EMBEDDING_MODEL instead of silently falling back to the default model. --- crates/agentflare-store/src/embed.rs | 23 +++++++++++--- .../src/embedding_pipeline/download.rs | 31 ++++++++++++++----- .../src/embedding_pipeline/model_registry.rs | 29 ++++++++--------- 3 files changed, 56 insertions(+), 27 deletions(-) diff --git a/crates/agentflare-store/src/embed.rs b/crates/agentflare-store/src/embed.rs index 8d5dcbea..10953823 100644 --- a/crates/agentflare-store/src/embed.rs +++ b/crates/agentflare-store/src/embed.rs @@ -1,9 +1,15 @@ -pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { +/// Returns `None` if `a` and `b` have different lengths — `zip` would +/// otherwise silently ignore the excess coordinates and rank on a partial +/// vector instead of erroring. +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Option { + if a.len() != b.len() { + return None; + } let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); let na: f32 = a.iter().map(|x| x * x).sum(); let nb: f32 = b.iter().map(|x| x * x).sum(); let denom = na.sqrt() * nb.sqrt(); - if denom < 1e-12 { 0.0 } else { dot / denom } + Some(if denom < 1e-12 { 0.0 } else { dot / denom }) } pub fn normalize(v: &mut [f32]) { @@ -23,21 +29,28 @@ mod tests { fn cosine_similarity_identical() { let a = vec![1.0, 0.0, 0.0]; let b = vec![1.0, 0.0, 0.0]; - assert!((cosine_similarity(&a, &b) - 1.0).abs() < 1e-6); + assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-6); } #[test] fn cosine_similarity_orthogonal() { let a = vec![1.0, 0.0]; let b = vec![0.0, 1.0]; - assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6); + assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 1e-6); } #[test] fn cosine_similarity_zero_denom() { let a = vec![0.0, 0.0]; let b = vec![1.0, 0.0]; - assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6); + assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 1e-6); + } + + #[test] + fn cosine_similarity_mismatched_lengths() { + let a = vec![1.0, 0.0]; + let b = vec![1.0, 0.0, 0.0]; + assert_eq!(cosine_similarity(&a, &b), None); } #[test] diff --git a/crates/agentflare-store/src/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs index 06e33f0e..751362da 100644 --- a/crates/agentflare-store/src/embedding_pipeline/download.rs +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -21,7 +21,7 @@ pub fn ensure_model(model_dir: &Path, config: &ModelConfig) -> anyhow::Result anyhow::Result anyhow::Result anyhow::Result { Ok(format!("{result:x}")) } -fn read_lockfile(model_dir: &Path) -> BTreeMap { - std::fs::read_to_string(model_dir.join(LOCKFILE)) - .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or_default() +/// A missing lockfile is expected on first run (`Ok(empty)`); an existing but +/// unreadable or malformed one is not — silently treating it as empty would +/// make every present file look unpinned and skip verification entirely. +fn read_lockfile(model_dir: &Path) -> anyhow::Result> { + let path = model_dir.join(LOCKFILE); + match std::fs::read_to_string(&path) { + Ok(s) => serde_json::from_str(&s) + .map_err(|e| anyhow::anyhow!("Malformed lockfile {}: {e}", path.display())), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(BTreeMap::new()), + Err(e) => Err(anyhow::anyhow!( + "Cannot read lockfile {}: {e}", + path.display() + )), + } } fn write_lockfile(model_dir: &Path, lock: &BTreeMap) -> anyhow::Result<()> { diff --git a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs index 319ac5e3..b61f5ef9 100644 --- a/crates/agentflare-store/src/embedding_pipeline/model_registry.rs +++ b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs @@ -49,16 +49,15 @@ impl CustomModelSpec { _ => '-', }); } - if let Some(rev) = &self.revision { - slug.push('-'); - for c in rev.chars().take(16) { - slug.push(match c { - 'a'..='z' | '0'..='9' | '-' => c, - 'A'..='Z' => c.to_ascii_lowercase(), - _ => '-', - }); - } - } + // Hash repo+revision instead of truncating the revision to 16 chars: + // two different (repo, revision) pairs that happen to share a + // 16-char revision prefix would otherwise collide on the same cache + // directory. + let digest = blake3::hash( + format!("{}@{}", self.repo, self.revision.as_deref().unwrap_or("")).as_bytes(), + ); + slug.push('-'); + slug.push_str(&digest.to_hex()[..16]); slug } } @@ -208,8 +207,10 @@ impl ModelConfig { } } -pub fn resolve_model() -> EmbeddingModel { - let env_val = std::env::var("AGENTFLARE_EMBEDDING_MODEL").ok(); - let name = env_val.as_deref().unwrap_or("default"); - EmbeddingModel::from_str_name(name).unwrap_or(EmbeddingModel::DEFAULT) +pub fn resolve_model() -> anyhow::Result { + match std::env::var("AGENTFLARE_EMBEDDING_MODEL") { + Ok(name) => EmbeddingModel::from_str_name(&name) + .ok_or_else(|| anyhow::anyhow!("invalid AGENTFLARE_EMBEDDING_MODEL: {name:?}")), + Err(_) => Ok(EmbeddingModel::DEFAULT), + } } From acb01229828e9f042070b397a438c141ab71a4ac Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 12:28:31 +0530 Subject: [PATCH 14/16] fix(agentflare-store): collision-resistant model slug, name-based ONNX I/O, correct WordPiece fallback - CustomModelSpec::storage_slug: hash repo+revision into the cache dir name instead of truncating the revision to 16 chars, which let two different (repo, revision) pairs collide on the same cache directory. - EmbeddingEngine::load_model: resolve ONNX input/output tensors by declared name (with a positional fallback for the two inputs every model has) instead of blind index 0/1/2, which breaks if a model export ever reorders its inputs/outputs. - wordpiece_encode: an unsegmentable word now collapses to a single [UNK], matching standard WordPiece, instead of one UNK per unmatched character interleaved with whatever subwords did match. - vec_search_ranks_by_similarity test: use directionally distinct embeddings; the old ones were collinear and scored identically regardless of magnitude, so the test wasn't exercising rank order. --- .../src/embedding_pipeline/mod.rs | 93 ++++++++++++++----- .../src/embedding_pipeline/tokenizer.rs | 29 +++++- 2 files changed, 96 insertions(+), 26 deletions(-) diff --git a/crates/agentflare-store/src/embedding_pipeline/mod.rs b/crates/agentflare-store/src/embedding_pipeline/mod.rs index 977527d5..6aa952d0 100644 --- a/crates/agentflare-store/src/embedding_pipeline/mod.rs +++ b/crates/agentflare-store/src/embedding_pipeline/mod.rs @@ -33,7 +33,7 @@ struct InputNames { impl EmbeddingEngine { pub fn load(model_dir: &Path) -> anyhow::Result { - let selected = model_registry::resolve_model(); + let selected = model_registry::resolve_model()?; Self::load_model(model_dir, selected) } @@ -60,33 +60,52 @@ impl EmbeddingEngine { .iter() .map(|i| i.name().to_string()) .collect(); + let output_names_list: Vec = session + .outputs() + .iter() + .map(|o| o.name().to_string()) + .collect(); + + // Resolved by name first — ONNX export order isn't guaranteed + // stable across model versions — with a positional fallback only + // for the two inputs every embedding model has. + let input_ids_name = find_input_name(&input_names_list, "input_ids") + .or_else(|| input_names_list.first().cloned()) + .ok_or_else(|| anyhow::anyhow!("Model {} has no inputs", config.name))?; + let attention_mask_name = find_input_name(&input_names_list, "attention_mask") + .or_else(|| input_names_list.get(1).cloned()) + .ok_or_else(|| { + anyhow::anyhow!("Model {} is missing an attention_mask input", config.name) + })?; + let named_token_type_ids = find_input_name(&input_names_list, "token_type_ids"); let token_type_ids = if config.needs_token_type_ids { - if input_names_list.len() < 3 { - anyhow::bail!( - "Model {} requires token_type_ids but only has {} inputs", - config.name, - input_names_list.len() - ); - } - Some(input_names_list[2].clone()) - } else if input_names_list.len() >= 3 { - Some(input_names_list[2].clone()) + Some( + named_token_type_ids + .clone() + .or_else(|| input_names_list.get(2).cloned()) + .ok_or_else(|| { + anyhow::anyhow!( + "Model {} requires token_type_ids but only has {} inputs", + config.name, + input_names_list.len() + ) + })?, + ) } else { - None + named_token_type_ids }; - let output_name = session - .outputs() - .first() - .map(|o| o.name().to_string()) + let output_name = find_output_name(&output_names_list) + .or_else(|| output_names_list.first().cloned()) .ok_or_else(|| anyhow::anyhow!("Model has no named outputs"))?; let dimensions = detect_dimensions( &config, &model_path, &tokenizer, - &input_names_list, + &input_ids_name, + &attention_mask_name, &token_type_ids, &output_name, )?; @@ -99,8 +118,8 @@ impl EmbeddingEngine { model_id, model_config: config, input_names: InputNames { - input_ids: input_names_list[0].clone(), - attention_mask: input_names_list[1].clone(), + input_ids: input_ids_name, + attention_mask: attention_mask_name, token_type_ids, }, output_name, @@ -236,11 +255,37 @@ fn tokenize(tokenizer: &TokenizerKind, text: &str, max_len: usize) -> TokenizedI } } +/// Finds an input/output whose ONNX-declared name matches `keyword` (exact +/// match preferred, substring as a fallback) — case-insensitive since +/// exporters aren't consistent about casing. +fn find_input_name(names: &[String], keyword: &str) -> Option { + names + .iter() + .find(|n| n.eq_ignore_ascii_case(keyword)) + .or_else(|| { + names + .iter() + .find(|n| n.to_ascii_lowercase().contains(keyword)) + }) + .cloned() +} + +fn find_output_name(names: &[String]) -> Option { + const PREFERRED: &[&str] = &[ + "last_hidden_state", + "sentence_embedding", + "hidden_state", + "embedding", + ]; + PREFERRED.iter().find_map(|kw| find_input_name(names, kw)) +} + fn detect_dimensions( config: &ModelConfig, model_path: &Path, tokenizer: &TokenizerKind, - input_names: &[String], + input_ids_name: &str, + attention_mask_name: &str, token_type_ids: &Option, output_name: &str, ) -> anyhow::Result { @@ -269,14 +314,14 @@ fn detect_dimensions( let type_array = ndarray::Array2::from_shape_vec((1, seq_len), type_vec)?; let type_tensor = ort::value::Tensor::from_array(type_array)?; session.run(ort::inputs![ - input_names[0].as_str() => ids_tensor, - input_names[1].as_str() => mask_tensor, + input_ids_name => ids_tensor, + attention_mask_name => mask_tensor, type_id.as_str() => type_tensor, ])? } else { session.run(ort::inputs![ - input_names[0].as_str() => ids_tensor, - input_names[1].as_str() => mask_tensor, + input_ids_name => ids_tensor, + attention_mask_name => mask_tensor, ])? }; diff --git a/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs index 1cdc4b50..f0619537 100644 --- a/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs +++ b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs @@ -166,8 +166,10 @@ impl WordPieceTokenizer { end -= 1; } if !matched { - tokens.push(self.unk_id); - start += 1; + // Standard WordPiece: a word that can't be fully segmented + // maps to a single [UNK], not one UNK per unmatched + // character interleaved with whatever subwords did match. + return vec![self.unk_id]; } } tokens @@ -310,6 +312,18 @@ pub enum HfTokenizerInner { Bpe(BpeTokenizer), } +/// Parses only `model.type`/`model.vocab`/`model.merges` from a HF +/// `tokenizer.json`; pre-tokenization always runs this crate's own +/// whitespace/BERT-punctuation splitter + lowercasing (see +/// [`WordPieceTokenizer::pre_tokenize`]), not the file's own `normalizer`, +/// `pre_tokenizer`, or `post_processor` sections (NFD/accent-stripping, +/// byte-level/Metaspace pre-tokenizers, template special-token insertion, +/// etc). This matches the built-in MiniLM/Nomic models, which use plain +/// BERT-style WordPiece, but a custom `hf:owner/repo` model that relies on a +/// non-default normalizer or pre-tokenizer will tokenize differently than +/// the reference HF implementation. Full parity would mean implementing the +/// whole tokenizers normalizer/pre-tokenizer/post-processor grammar, which +/// is out of scope here — known limitation, not a bug in the common path. pub struct HfTokenizerWrapper { inner: HfTokenizerInner, } @@ -478,6 +492,17 @@ mod tests { assert_eq!(out.input_ids, vec![2, 4, 3]); } + #[test] + fn wordpiece_unsegmentable_word_is_single_unk() { + let tok = HfTokenizerWrapper::from_json(WORDPICE_JSON).unwrap(); + // "helloz" greedily matches "hello" then fails on the trailing "z" + // (no "##z" in vocab) — the whole word should collapse to one + // [UNK], not "hello" followed by a stray [UNK]. + let out = tok.encode("helloz", 32); + // [CLS]=2, [UNK]=1, [SEP]=3 + assert_eq!(out.input_ids, vec![2, 1, 3]); + } + #[test] fn unknown_model_type_fails_loudly() { let json = r#"{ "model": { "type": "Unigram", "vocab": {} } }"#; From fd4e92cbedf9ecac6bf7d7ce4638b9f4a1f5ecce Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 12:29:02 +0530 Subject: [PATCH 15/16] fix(flare-output): treat blank model env vars as unset, keep pre-rename backups discoverable - call_via_api: a blank FLARE_OUTPUT_MODEL/CAVEMAN_MODEL value no longer wins over the next fallback in the chain. - backup_path_for (OutOfTree): the backup namespace directory was renamed from 'caveman' to 'flare-output' with a hard rename and no fallback, so a user's pre-rename backups became invisible to both the 'backup already exists' guard and Report::original_path. Now resolves to an existing legacy path if the new one doesn't have one yet. --- crates/flare-output/src/compress.rs | 62 +++++++++++++++++++++++++---- crates/flare-output/src/llm.rs | 13 ++++-- 2 files changed, 65 insertions(+), 10 deletions(-) diff --git a/crates/flare-output/src/compress.rs b/crates/flare-output/src/compress.rs index 93dd3907..9239eb5c 100644 --- a/crates/flare-output/src/compress.rs +++ b/crates/flare-output/src/compress.rs @@ -131,6 +131,14 @@ pub fn compress( }) } +fn out_of_tree_backup_dir(namespace: &str) -> PathBuf { + dirs::cache_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("agentflare") + .join(namespace) + .join("backups") +} + fn backup_path_for(target: &Path, mode: BackupMode) -> PathBuf { match mode { BackupMode::Sibling => { @@ -141,11 +149,6 @@ fn backup_path_for(target: &Path, mode: BackupMode) -> PathBuf { target.with_file_name(format!("{file_name}.orig")) } BackupMode::OutOfTree => { - let base = dirs::cache_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join("agentflare") - .join("flare-output") - .join("backups"); // Hash the full parent path, not just its last component — two // files with the same name under differently-located but // identically-named parent dirs (e.g. "project-a/docs/README.md" @@ -163,8 +166,23 @@ fn backup_path_for(target: &Path, mode: BackupMode) -> PathBuf { .file_stem() .map(|s| s.to_string_lossy().to_string()) .unwrap_or_default(); - base.join(format!("{dir_hash:016x}")) - .join(format!("{stem}.original.md")) + let sub = format!("{dir_hash:016x}"); + let file = format!("{stem}.original.md"); + + let current = out_of_tree_backup_dir("flare-output") + .join(&sub) + .join(&file); + // The backup namespace was renamed from "caveman" to + // "flare-output". If a pre-rename backup already sits at the + // old path, keep resolving to it — otherwise it (and the + // "backup already exists" guard above, which checks this same + // path) would silently stop seeing it. + let legacy = out_of_tree_backup_dir("caveman").join(&sub).join(&file); + if !current.exists() && legacy.exists() { + legacy + } else { + current + } } } } @@ -326,6 +344,36 @@ mod tests { ); } + #[test] + fn out_of_tree_backup_falls_back_to_pre_rename_caveman_path() { + let dir = tempdir().unwrap(); + let target = write(dir.path(), "legacy-fallback-test.md", "content"); + + // Nothing backed up anywhere yet: resolves under the current + // "flare-output" namespace. + let resolved = backup_path_for(&target, BackupMode::OutOfTree); + assert!(resolved.to_string_lossy().contains("flare-output")); + assert!(!resolved.exists()); + + // Simulate a backup left behind by a pre-rename version of this + // tool, at the equivalent path under the old "caveman" namespace. + let legacy = PathBuf::from( + resolved + .to_string_lossy() + .replace("flare-output", "caveman"), + ); + std::fs::create_dir_all(legacy.parent().unwrap()).unwrap(); + std::fs::write(&legacy, "pre-rename backup").unwrap(); + + let resolved_again = backup_path_for(&target, BackupMode::OutOfTree); + assert_eq!( + resolved_again, legacy, + "must resolve to the existing legacy backup instead of orphaning it" + ); + + std::fs::remove_file(&legacy).unwrap(); + } + #[test] fn shadow_copy_write_needs_no_backup() { let dir = tempdir().unwrap(); diff --git a/crates/flare-output/src/llm.rs b/crates/flare-output/src/llm.rs index a5013d11..cf0c3c34 100644 --- a/crates/flare-output/src/llm.rs +++ b/crates/flare-output/src/llm.rs @@ -25,10 +25,17 @@ impl Llm for RealLlm { } } +fn env_model(key: &str) -> Option { + // A blank value (e.g. `FLARE_OUTPUT_MODEL=""` from an unset shell var + // interpolated into an env file) must not win over the next fallback — + // treat it the same as unset. + std::env::var(key).ok().filter(|v| !v.trim().is_empty()) +} + fn call_via_api(api_key: &str, prompt: &str) -> Result { - let model = std::env::var("FLARE_OUTPUT_MODEL") - .or_else(|_| std::env::var("CAVEMAN_MODEL")) - .unwrap_or_else(|_| "claude-sonnet-4-5".to_string()); + let model = env_model("FLARE_OUTPUT_MODEL") + .or_else(|| env_model("CAVEMAN_MODEL")) + .unwrap_or_else(|| "claude-sonnet-4-5".to_string()); let body = serde_json::json!({ "model": model, "max_tokens": 8192, From 2f0ab6f4cd41143f22ba4fdecc10fc6c05502283 Mon Sep 17 00:00:00 2001 From: Shivakumar Date: Sat, 18 Jul 2026 13:27:26 +0530 Subject: [PATCH 16/16] fix(flare-output): build legacy caveman test path by component, not string replace --- crates/flare-output/src/compress.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/flare-output/src/compress.rs b/crates/flare-output/src/compress.rs index 9239eb5c..923bc3eb 100644 --- a/crates/flare-output/src/compress.rs +++ b/crates/flare-output/src/compress.rs @@ -357,11 +357,9 @@ mod tests { // Simulate a backup left behind by a pre-rename version of this // tool, at the equivalent path under the old "caveman" namespace. - let legacy = PathBuf::from( - resolved - .to_string_lossy() - .replace("flare-output", "caveman"), - ); + let current_root = out_of_tree_backup_dir("flare-output"); + let relative = resolved.strip_prefix(¤t_root).unwrap(); + let legacy = out_of_tree_backup_dir("caveman").join(relative); std::fs::create_dir_all(legacy.parent().unwrap()).unwrap(); std::fs::write(&legacy, "pre-rename backup").unwrap();