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/Cargo.lock b/Cargo.lock index 22d235ca..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]] @@ -221,6 +222,30 @@ dependencies = [ "thiserror", ] +[[package]] +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]] name = "aho-corasick" version = "1.1.4" @@ -289,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" @@ -298,6 +329,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 +443,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 +481,32 @@ 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 = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.12.0" @@ -599,6 +682,22 @@ 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" +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" @@ -632,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" @@ -692,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" @@ -853,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" @@ -1047,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" @@ -1408,18 +1557,43 @@ 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" 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" @@ -1453,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" @@ -1474,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" @@ -1529,18 +1753,108 @@ 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" 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" @@ -1570,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" @@ -1602,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 = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] [[package]] name = "potential_utf" @@ -1719,6 +2051,41 @@ 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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_users" version = "0.5.2" @@ -2008,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" @@ -2034,6 +2410,35 @@ dependencies = [ "syn", ] +[[package]] +name = "scopeguard" +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" @@ -2186,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" @@ -2627,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" @@ -2639,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" @@ -2783,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 0e91839e..71be4dda 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] @@ -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 new file mode 100644 index 00000000..1aa58f72 --- /dev/null +++ b/crates/agentflare-store/Cargo.toml @@ -0,0 +1,34 @@ +[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"] } +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 new file mode 100644 index 00000000..e5287d19 --- /dev/null +++ b/crates/agentflare-store/src/blobs.rs @@ -0,0 +1,265 @@ +use crate::Store; +use rusqlite::{OptionalExtension, params}; +use serde::Serialize; +use std::path::Path; + +#[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; + +fn blob_disk_path(root: &Path, hash: &str) -> PathBuf { + root.join("blobs").join(&hash[..2]).join(hash) +} + +/// `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); + 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> { + let path = blob_disk_path(root, hash); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(&path, data)?; + Ok(()) +} + +fn delete_disk_blob(root: &Path, hash: &str) { + let path = blob_disk_path(root, hash); + let _ = std::fs::remove_file(&path); +} + +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(); + let is_memory = self.is_memory(); + + // 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], + |_| Ok(()), + ) + .optional()? + .is_some(); + + if exists { + tx.execute( + "UPDATE store_blobs SET ref_count = ref_count + 1 WHERE hash = ?1", + params![hash], + )?; + tx.commit()?; + return Ok(hash); + } + + 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() { + tx.execute( + "INSERT INTO store_blob_chunks (hash, chunk_index, data) VALUES (?1, ?2, ?3)", + params![hash, i as i64, chunk], + )?; + } + } + + 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) + } + + 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], + |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), + }; + + if !self.is_memory() { + return read_disk_blob(&self.root, hash) + .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e))); + } + + 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(); + + 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 conn = self.conn(); + let n = 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 conn = self.conn(); + 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 { + removed = tx + .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 { + 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) + } +} + +#[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()); + } + + #[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 new file mode 100644 index 00000000..ea55e0f1 --- /dev/null +++ b/crates/agentflare-store/src/documents.rs @@ -0,0 +1,751 @@ +use crate::Store; +use rusqlite::{OptionalExtension, params}; +use serde::{Deserialize, 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, + pub project_id: String, + pub path: String, + pub snippet: String, + 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( + conn: &rusqlite::Connection, + row_id: i64, + content: &str, + ) -> rusqlite::Result<()> { + conn.execute( + "DELETE FROM store_docs_fts WHERE rowid = ?1", + params![row_id], + )?; + 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(); + + // 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", + params![project_id, path], + |row| { + Ok(( + row.get::<_, String>(0)?, + 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, old_blob_hash, old_mime)) = + existing + { + let new_version = old_version + 1; + let history_id = db_kit::ids::new_id(); + + // Snapshot current version to history + 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], + )?; + + tx.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 { + tx.execute( + "UPDATE store_documents SET title = ?1 WHERE id = ?2", + params![title, existing_id], + )?; + } + if let Some(doc_type) = &opts.doc_type { + tx.execute( + "UPDATE store_documents SET doc_type = ?1 WHERE id = ?2", + params![doc_type, existing_id], + )?; + } + if opts.blob_hash.is_some() { + tx.execute( + "UPDATE store_documents SET blob_hash = ?1 WHERE id = ?2", + params![opts.blob_hash, existing_id], + )?; + } + if let Some(mime) = &opts.mime { + 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()); + tx.execute( + "UPDATE store_documents SET tags = ?1 WHERE id = ?2", + params![json, existing_id], + )?; + } + if opts.session_id.is_some() { + tx.execute( + "UPDATE store_documents SET session_id = ?1 WHERE id = ?2", + params![opts.session_id, existing_id], + )?; + } + if let Some(source) = &opts.source { + tx.execute( + "UPDATE store_documents SET source = ?1 WHERE id = ?2", + params![source, existing_id], + )?; + } + + Self::doc_sync_fts(&tx, rowid, content)?; + tx.commit()?; + drop(conn); + self.doc_get(&existing_id).map(|o| o.unwrap()) + } else { + 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(); + + // 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)", + params![ + id, project_id, path, content, title, doc_type, opts.blob_hash, + mime, tags_json, opts.session_id, source, now + ], + )?; + let rowid = tx.last_insert_rowid(); + Self::doc_sync_fts(&tx, rowid, content)?; + tx.commit()?; + 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, + }) + } + } + + 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, + session_id, source, version, created_at, updated_at, deleted_at + FROM store_documents WHERE id = ?1", + params![id], + 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) = conn + .query_row( + "SELECT rowid FROM store_documents WHERE id = ?1", + params![id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + conn.execute( + "UPDATE store_documents SET deleted_at = ?1 WHERE id = ?2", + params![now, id], + )?; + 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 { + let conn = self.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()? + 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> { + 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 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 + 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 conn = self.conn(); + let now = db_kit::ids::now(); + let bytes: &[u8] = bytemuck::cast_slice(embedding); + 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], + )?; + Ok(n > 0) + } + + 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() + } + + pub fn doc_vec_search( + &self, + project_id: &str, + query_vec: &[f32], + limit: usize, + ) -> rusqlite::Result> { + 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 + WHERE d.project_id = ?1 AND d.deleted_at IS NULL", + )?; + 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)?; + let path: String = row.get(2)?; + let blob: Vec = row.get(3)?; + Ok((id, project_id, path, blob)) + })? + .collect::>>()?; + let mut results: Vec<(f64, DocMatch)> = rows + .into_iter() + .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 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], Self::row_to_document)?; + 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(); + + // 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, 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] + 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); + } + + #[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"); + } + + #[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/embed.rs b/crates/agentflare-store/src/embed.rs new file mode 100644 index 00000000..10953823 --- /dev/null +++ b/crates/agentflare-store/src/embed.rs @@ -0,0 +1,63 @@ +/// 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(); + Some(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).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).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).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] + 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/embedding_pipeline/download.rs b/crates/agentflare-store/src/embedding_pipeline/download.rs new file mode 100644 index 00000000..751362da --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/download.rs @@ -0,0 +1,364 @@ +use std::collections::BTreeMap; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use super::model_registry::ModelConfig; + +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); + + // 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 + ); + std::fs::remove_file(&local_path).map_err(|e| { + anyhow::anyhow!( + "Failed to remove corrupt model file {}: {e}. Refusing to \ + continue with unverified data left on disk.", + local_path.display() + ) + })?; + 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 or invalid, 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(()) +} + +/// 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)?; + 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}")) +} + +/// 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<()> { + 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(()) +} + +#[cfg(all(test, feature = "embeddings"))] +mod tests { + use super::*; + use crate::embedding_pipeline::model_registry::{EmbeddingModel, ModelConfig}; + use std::io::{Read, Write}; + use std::net::TcpListener; + 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"; + + /// 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/mod.rs b/crates/agentflare-store/src/embedding_pipeline/mod.rs new file mode 100644 index 00000000..6aa952d0 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/mod.rs @@ -0,0 +1,334 @@ +pub mod download; +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::{HfTokenizerWrapper, TokenizedInput, WordPieceTokenizer}; + +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, +} + +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 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 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 { + 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 { + named_token_type_ids + }; + + 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_ids_name, + &attention_mask_name, + &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_ids_name, + attention_mask: attention_mask_name, + 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)?; + + hidden = pooling::mean_pool( + &hidden, + &input.attention_mask, + input.input_ids.len(), + self.dimensions, + ); + + embed::normalize(&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; + + embed::normalize(&mut hidden); + Ok(hidden) + } + + pub fn model_id(&self) -> &EmbeddingModel { + &self.model_id + } + + 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); + } + dirs::cache_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("agentflare") + .join("models") + } + + 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), + } +} + +/// 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_ids_name: &str, + attention_mask_name: &str, + 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(model_path) + .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_ids_name => ids_tensor, + attention_mask_name => mask_tensor, + type_id.as_str() => type_tensor, + ])? + } else { + session.run(ort::inputs![ + input_ids_name => ids_tensor, + attention_mask_name => 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..b61f5ef9 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/model_registry.rs @@ -0,0 +1,216 @@ +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(), + _ => '-', + }); + } + // 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 + } +} + +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, + base_url_override: None, + }, + 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, + base_url_override: None, + }, + 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, + base_url_override: None, + }, + } + } + + 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, + /// 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, + 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 { + if let Some(base) = &self.base_url_override { + return base.clone(); + } + 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() -> 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), + } +} 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..4db7238f --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/pooling.rs @@ -0,0 +1,28 @@ +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, s) in sum.iter_mut().enumerate() { + if let Some(&val) = hidden_states.get(offset + d) { + *s += val; + } + } + count += 1.0; + } + } + + if count > 0.0 { + for val in &mut sum { + *val /= count; + } + } + sum +} 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..f0619537 --- /dev/null +++ b/crates/agentflare-store/src/embedding_pipeline/tokenizer.rs @@ -0,0 +1,522 @@ +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 { + // 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 + } +} + +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), +} + +/// 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, +} + +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 model = parsed + .get("model") + .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 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"))?; + + 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 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." + ), + }; + + Ok(Self { inner }) + } + + pub fn encode(&self, text: &str, max_len: usize) -> TokenizedInput { + 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, + '!' | '"' + | '#' + | '$' + | '%' + | '&' + | '\'' + | '(' + | ')' + | '*' + | '+' + | ',' + | '-' + | '.' + | '/' + | ':' + | ';' + | '<' + | '=' + | '>' + | '?' + | '@' + | '[' + | '\\' + | ']' + | '^' + | '_' + | '`' + | '{' + | '|' + | '}' + | '~' + ) + } 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 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": {} } }"#; + 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()); + } +} diff --git a/crates/agentflare-store/src/kv.rs b/crates/agentflare-store/src/kv.rs new file mode 100644 index 00000000..c2b08899 --- /dev/null +++ b/crates/agentflare-store/src/kv.rs @@ -0,0 +1,129 @@ +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 conn = self.conn(); + let now = db_kit::ids::now(); + 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> { + 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() + } + + 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()) + } + + pub fn kv_delete(&self, key: &str) -> rusqlite::Result { + 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 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}'); + 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..61392244 --- /dev/null +++ b/crates/agentflare-store/src/lib.rs @@ -0,0 +1,80 @@ +pub mod blobs; +pub mod documents; +pub mod embed; +pub mod kv; +pub mod leases; +pub mod migrate; +pub mod migrations; + +#[cfg(feature = "embeddings")] +pub mod embedding_pipeline; + +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: parking_lot::Mutex, + 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: parking_lot::Mutex::new(conn), + root, + }) + } + + pub fn open_memory() -> Result { + let conn = db_kit::open_memory(&migrations::migrations())?; + Ok(Self { + conn: parking_lot::Mutex::new(conn), + root: PathBuf::from(":memory:"), + }) + } + + pub fn conn(&self) -> parking_lot::MutexGuard<'_, Connection> { + self.conn.lock() + } + + 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..e2aaa4cc --- /dev/null +++ b/crates/agentflare-store/src/migrate.rs @@ -0,0 +1,108 @@ +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)?; + + // 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 { + upsert(key, &serde_json::to_vec(value)?)?; + } + upsert(MIGRATION_MARKER, &serde_json::to_vec(&now)?)?; + + tx.commit()?; + 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..032b6334 --- /dev/null +++ b/crates/agentflare-store/src/migrations.rs @@ -0,0 +1,84 @@ +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 + );", + ), + 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/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-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/compress.rs b/crates/flare-output/src/compress.rs index d97381d7..923bc3eb 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("caveman") - .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,34 @@ 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 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(); + + 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 ccfcbd2d..cf0c3c34 100644 --- a/crates/flare-output/src/llm.rs +++ b/crates/flare-output/src/llm.rs @@ -25,8 +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("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, 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!( 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/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 new file mode 100644 index 00000000..bcc8b427 --- /dev/null +++ b/src/store.rs @@ -0,0 +1,49 @@ +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("store.db") +} + +/// 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)] +mod tests { + use super::*; + use crate::paths::test_support::with_temp_home; + + #[test] + fn open_and_close() { + with_temp_home(|| { + 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()); + }); + } +}