diff --git a/CHANGELOG.md b/CHANGELOG.md index 188df5528..e212099cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `embedding_query_prefix` / `embedding_document_prefix` config keys (env: + `AI_MEMORY_EMBEDDING_QUERY_PREFIX` / `AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX`) + for the `openai` and `openai-compat` embedders: an optional string + prepended to query / document text before the existing truncation, so + truncation still bounds the whole input. Asymmetric embedding models need + a query-side instruction their publisher specifies; the OpenAI-compatible + `/v1/embeddings` wire format has no field for it. + `nvidia/Nemotron-3-Embed-1B-BF16` and base E5 models use a simple + `"query: "` / `"passage: "` pair; instruction-tuned E5 variants and + Qwen3-Embedding instead need a task-instruction string on the query side + only (documents stay plain). Empty by default — no behaviour change when + unset, and not trimmed (nor is a present-but-empty env-var override, which + now clears a `config.toml` value), so a publisher's trailing space is + preserved. A non-empty `embedding_document_prefix` is folded into the + stored embedding identity (`Embedder::model_identity`), so a document + prefix change makes retrieval, backfill, and cleanup treat existing pages + as stale and re-embed them automatically — an empty prefix keeps the + pre-existing identity, so upgrading installs need no migration, and a + query-only prefix change never needs a rebuild. See + `docs/llm-providers.md`. + ### Fixed +- `memory_query`'s vector stream called the generic `Embedder::embed` + instead of `embed_query` on the configured embedder, so a + query/document-asymmetric embedder (Google's task-typed embeddings, or + the new query/document prefixes above) embedded the search query on the + document side instead of the query side. - `companions/ai-memory-macos/build.sh` no longer fails on machines whose active developer directory is Command Line Tools only: SwiftUI `@State` needs the `SwiftUIMacros` plugin shipped with full Xcode, so the script diff --git a/bin/ai-memory b/bin/ai-memory index 18b6dcc12..9cd096824 100755 --- a/bin/ai-memory +++ b/bin/ai-memory @@ -612,6 +612,18 @@ do fi done +# These two are presence-based, not non-empty-based like the loop above: an +# operator sets one to the empty string to deliberately clear a +# config.toml-configured prefix without editing the file (see +# Config::load's figment overlay), and a present-but-empty value must reach +# the container for that to work — `[ -n ]` above would drop it, making the +# wrapper indistinguishable from the var never having been set at all. +for var in AI_MEMORY_EMBEDDING_QUERY_PREFIX AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX; do + if [ -n "${!var+x}" ]; then + ENV_ARGS+=(-e "${var}") + fi +done + # The wrapper itself runs the CLI inside a short-lived helper container, while # the README server runs in the long-lived ai-memory container and publishes # 127.0.0.1:49374 on the host. Inside a normal bridge-network helper, diff --git a/bin/ai-memory.ps1 b/bin/ai-memory.ps1 index 34f9cbc4a..09c9ff9f5 100644 --- a/bin/ai-memory.ps1 +++ b/bin/ai-memory.ps1 @@ -174,6 +174,30 @@ foreach ($Name in @( } } +# Presence-based, not non-empty-based like the loop above: an operator sets +# one of these to the empty string to deliberately clear a +# config.toml-configured prefix without editing the file (see +# Config::load's figment overlay), and a present-but-empty value must reach +# the container for that to work — `IsNullOrEmpty` above would drop it, +# making the wrapper indistinguishable from the variable never having been +# set at all. `GetEnvironmentVariable` returns `$null` only when the +# variable is truly unset, and `""` when it is set-but-empty, so a `-ne +# $null` check is exactly the presence test needed here. +# +# An operator's own `$env:NAME = ''` additionally needs PowerShell 7.5+ +# (first built on .NET 9) to leave a set-but-empty variable rather than +# deleting it; not exercised on a real pwsh runtime. +# https://learn.microsoft.com/en-us/dotnet/api/system.environment.setenvironmentvariable +# https://learn.microsoft.com/en-us/powershell/scripting/whats-new/what-s-new-in-powershell-75 +foreach ($Name in @( + "AI_MEMORY_EMBEDDING_QUERY_PREFIX", + "AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX" +)) { + if ($null -ne [Environment]::GetEnvironmentVariable($Name)) { + $DockerArgs += @("-e", $Name) + } +} + # Docker Desktop gives Windows no host networking for Linux containers, so a # thin-client command (status, search, bootstrap, ...) reaches the loopback- # published server from this helper container through Docker Desktop's host diff --git a/crates/ai-memory-cli/src/commands/serve.rs b/crates/ai-memory-cli/src/commands/serve.rs index d2c1811b7..21448f1fe 100644 --- a/crates/ai-memory-cli/src/commands/serve.rs +++ b/crates/ai-memory-cli/src/commands/serve.rs @@ -1633,7 +1633,7 @@ async fn start_maintenance_scheduler( .as_ref() .map(|e| ai_memory_consolidate::EmbeddingCoord { provider: e.provider().to_string(), - model: e.model().to_string(), + model: e.model_identity(), dim: e.dim(), }); @@ -2367,11 +2367,15 @@ async fn configure_embedder( } Err(e) => return Err(e).context("building embedder from config"), }; + // Not `.model()`: the running triple must match what pages are + // actually stored under, which a configured document prefix changes. + // See `Embedder::model_identity`. + let configured_model_identity = embedder.model_identity(); let mismatch = store .reader .embedding_meta_for_mismatch( embedder.provider().into(), - embedder.model().into(), + configured_model_identity.clone(), embedder.dim(), ) .await?; @@ -2383,7 +2387,7 @@ async fn configure_embedder( tracing::warn!( stored = ?mismatch, configured_provider = embedder.provider(), - configured_model = embedder.model(), + configured_model = %configured_model_identity, configured_dim = embedder.dim(), "stored embeddings use a different (provider, model, dim) than configured; \ hybrid search ignores stale rows until pages are re-embedded — \ @@ -3999,7 +4003,7 @@ mod tests { ws, project, embedder.provider().to_string(), - embedder.model().to_string(), + embedder.model_identity(), embedder.dim(), ) .await diff --git a/crates/ai-memory-cli/src/config.rs b/crates/ai-memory-cli/src/config.rs index 1e3eae2e5..b544ecf95 100644 --- a/crates/ai-memory-cli/src/config.rs +++ b/crates/ai-memory-cli/src/config.rs @@ -441,6 +441,48 @@ pub struct Config { pub embedding_dim: Option, /// Optional embedding base URL override. pub embedding_base_url: Option, + /// Optional prefix prepended to every embedding **query** before it is + /// sent to the `openai` or `openai-compat` embedder, ahead of the + /// existing truncation. Unset (the default) is a no-op — no behaviour + /// change. Asymmetric self-hosted models need a query-side instruction + /// their publisher specifies; the OpenAI-compatible `/v1/embeddings` + /// wire format has no field for it, so the client prepends it instead. + /// `nvidia/Nemotron-3-Embed-1B-BF16` and base E5 models + /// (`intfloat/e5-base-v2`, multilingual E5, …) use a simple + /// `"query: "` / `"passage: "` pair (documents get + /// `embedding_document_prefix = "passage: "`). Instruction-tuned E5 + /// variants and Qwen3-Embedding instead need a full task-instruction + /// string on the query side only, with **different exact spacing each** + /// — leave `embedding_document_prefix` unset for both (their documents + /// are plain text, no prefix): + /// `e5-mistral-7b-instruct` wants + /// `"Instruct: {task description}\nQuery: "` (a trailing space after + /// `Query:`); Qwen3-Embedding wants + /// `"Instruct: {task description}\nQuery:"` (no trailing space — the + /// query text follows the colon directly). Not trimmed: a publisher's + /// trailing space or embedded newline is significant and preserved + /// verbatim. Ignored by `google` (which has its own built-in + /// query/document asymmetry), `voyage`, `local`, and `copilot`. + /// Changing only this key never requires re-embedding existing pages — + /// the query side has no stored identity. See `docs/llm-providers.md`. + /// Settable via `AI_MEMORY_EMBEDDING_QUERY_PREFIX` (figment's `Env` + /// provider would otherwise trim a trailing space; `Config::load` + /// overlays the raw env bytes for this key specifically). + pub embedding_query_prefix: Option, + /// Document-side counterpart of `embedding_query_prefix` (e.g. + /// `"passage: "` for Nemotron-3-Embed / base E5; see that field's doc + /// comment for which models this applies to). Unlike the query prefix, + /// this one IS folded into the stored embedding identity + /// (`Embedder::model_identity`): a non-empty value makes newly + /// embedded pages distinguishable from ones embedded before the + /// change (or under a different prefix), so `memory_query` and + /// `ai-memory embed`'s stale-row detection both treat a document-prefix + /// change like a model change — no manual `--force` needed, and an + /// empty value keeps the pre-existing (legacy) identity so upgrading + /// installs need no migration. Settable via + /// `AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX` (same raw-env overlay as + /// `embedding_query_prefix`). + pub embedding_document_prefix: Option, /// M8 retention-sweep parameters. The defaults give an ~80-day /// "survival floor" for unused episodic content (above the cold /// threshold), followed by ~180 days of tombstone grace before permanent @@ -875,6 +917,8 @@ impl Default for Config { embedding_model: None, embedding_dim: None, embedding_base_url: None, + embedding_query_prefix: None, + embedding_document_prefix: None, decay: DecaySettings::default(), maintenance: MaintenanceSettings::default(), dream: DreamSettings::default(), @@ -1316,6 +1360,19 @@ impl Config { figment = figment.merge(Toml::file(&resolved_config_path)); } figment = figment.merge(Env::prefixed("AI_MEMORY_").split("__")); + // The environment is read once, here, and passed down as data — + // never inside `overlay_embedding_prefixes` itself — so that + // function stays directly testable without mutating process env or + // cwd (see its doc comment). + figment = overlay_embedding_prefixes( + figment, + std::env::var("AI_MEMORY_EMBEDDING_QUERY_PREFIX") + .ok() + .as_deref(), + std::env::var("AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX") + .ok() + .as_deref(), + ); let mut config: Config = figment.extract().with_context(|| { format!( @@ -1881,6 +1938,14 @@ impl Config { } else { None }; + // Not `non_empty`: that trims, and a publisher's trailing space + // (e.g. Nemotron-3-Embed's `"query: "`) is significant. By the time + // `Load` has run, `self.embedding_query_prefix` already holds the + // exact configured bytes regardless of source (TOML or env) — see + // `Config::load`'s env-prefix overlay, which corrects for + // figment's `Env` provider trimming unquoted values. + let query_prefix = self.embedding_query_prefix.clone().unwrap_or_default(); + let document_prefix = self.embedding_document_prefix.clone().unwrap_or_default(); Ok(Some(EmbedderConfig { provider, model, @@ -1890,6 +1955,8 @@ impl Config { models_dir: Some(self.data_dir.join("models")), copilot_auth, defaulted, + query_prefix, + document_prefix, })) } @@ -2094,6 +2161,51 @@ fn env_string(name: &str) -> Option { }) } +/// Overlay the two embedding-prefix keys onto `figment` with their raw, +/// untrimmed values, whenever the corresponding parameter is `Some` (even +/// `Some("")` — an operator clearing a `config.toml`-set prefix back to +/// none via an empty env var; only `None`, the variable genuinely absent, +/// leaves a `config.toml` value or the default untouched). +/// +/// figment's `Env` provider parses each var's string as a loose value +/// (`figment::value::parse::value`), and its bare/unquoted branch calls +/// `.trim()` — so `AI_MEMORY_EMBEDDING_QUERY_PREFIX="query: "` would +/// otherwise reach `embedding_query_prefix` as `"query:"`, silently +/// dropping the publisher-significant trailing space (verified against +/// figment 0.10.19's vendored source, `src/value/parse.rs:78`). +/// [`Serialized`] values are handed to figment as already-typed data (via +/// `serde::Serialize`), so they never pass through that string parser and +/// so are never trimmed. Callers merge this after `Env::prefixed` so it +/// wins over the (possibly trimmed) value that provider already set. +/// +/// The values come in as parameters, already read by the caller, rather +/// than this function reading `std::env::var` itself — the same pattern +/// `ai-memory-cli/src/commands/path_util.rs`'s `agent_config_home` and +/// `ai-memory-hooks`'s `drain_with_live_token` use, and for the same +/// reason: it keeps this function directly unit-testable without +/// mutating process environment or the current directory. Both are +/// unsafe or actively harmful to do from a `#[test]` in this crate's +/// multi-threaded lib test binary — `std::env::set_var` is `unsafe` under +/// edition 2024 and forbidden workspace-wide, and even a "safe" wrapper +/// such as `figment::Jail` still calls `std::env::set_current_dir` on the +/// real process (verified against its vendored source, +/// `src/jail.rs:141`), racing every other test in the binary that reads +/// env or relies on cwd — e.g. `tests/suite/backfill_e2e.rs`'s +/// `Command::current_dir` calls. +fn overlay_embedding_prefixes( + mut figment: Figment, + query_prefix_env: Option<&str>, + document_prefix_env: Option<&str>, +) -> Figment { + if let Some(v) = query_prefix_env { + figment = figment.merge(Serialized::default("embedding_query_prefix", v)); + } + if let Some(v) = document_prefix_env { + figment = figment.merge(Serialized::default("embedding_document_prefix", v)); + } + figment +} + fn env_path(name: &str) -> Option { env_string(name).map(PathBuf::from) } @@ -3219,6 +3331,168 @@ mod tests { )); } + #[test] + fn embedding_prefixes_default_empty_and_are_not_trimmed_when_set() { + // Unset: EmbedderConfig carries empty strings, so downstream + // embedders see byte-identical behaviour to before this feature. + let unset = Config { + embedding_provider: Some("openai-compat".into()), + embedding_model: Some("nvidia/Nemotron-3-Embed-1B-BF16".into()), + embedding_dim: Some(2048), + embedding_base_url: Some("http://localhost:8000/v1".into()), + ..Config::default() + }; + let embedder = unset.embedder_config().unwrap().unwrap(); + assert_eq!(embedder.query_prefix, ""); + assert_eq!(embedder.document_prefix, ""); + + // Set: the publisher's exact strings pass through, including the + // significant trailing space — `non_empty`'s trim would corrupt it. + let set = Config { + embedding_query_prefix: Some("query: ".into()), + embedding_document_prefix: Some("passage: ".into()), + ..unset + }; + let embedder = set.embedder_config().unwrap().unwrap(); + assert_eq!(embedder.query_prefix, "query: "); + assert_eq!(embedder.document_prefix, "passage: "); + } + + /// Pure unit tests for `overlay_embedding_prefixes`: no process env or + /// cwd mutation anywhere here (the workspace forbids `std::env::set_var` + /// as `unsafe` under edition 2024, and `figment::Jail` calls + /// `std::env::set_current_dir` on the real process internally, racing + /// every other test in this multi-threaded lib test binary that + /// relies on cwd, such as `tests/suite/backfill_e2e.rs`'s + /// `Command::current_dir` calls). Each test builds its own minimal + /// `Figment` in memory instead, exactly mirroring what `Config::load` + /// does (`Serialized::defaults` as the base, optionally a lower-priority + /// `Serialized` merge standing in for a `config.toml` value), and + /// extracts a `Config` to assert on — the identical merge machinery the + /// real loader uses, with the "env value" supplied as a parameter + /// instead of read from the process. + #[test] + fn overlay_embedding_prefixes_preserves_trailing_whitespace() { + // The regression this guards: figment's `Env` provider parses an + // unquoted value with its loose-value parser, whose bare-value + // branch calls `.trim()` — so without this overlay a real + // `AI_MEMORY_EMBEDDING_QUERY_PREFIX="query: "` would arrive as + // `"query:"`, silently dropping the space the model publisher + // requires. `Serialized` bypasses that parser entirely. + let base = Figment::from(Serialized::defaults(Config::default())); + let overlaid = overlay_embedding_prefixes(base, Some("query: "), Some("passage: ")); + let cfg: Config = overlaid.extract().unwrap(); + assert_eq!(cfg.embedding_query_prefix.as_deref(), Some("query: ")); + assert_eq!(cfg.embedding_document_prefix.as_deref(), Some("passage: ")); + } + + #[test] + fn overlay_embedding_prefixes_none_leaves_a_lower_layer_untouched() { + // Simulates a `config.toml` value already merged in at lower + // priority; passing `None` (the env var genuinely absent) must not + // disturb it. + let base = Figment::from(Serialized::defaults(Config::default())).merge( + Serialized::default("embedding_query_prefix", "toml-query: "), + ); + let overlaid = overlay_embedding_prefixes(base, None, None); + let cfg: Config = overlaid.extract().unwrap(); + assert_eq!(cfg.embedding_query_prefix.as_deref(), Some("toml-query: ")); + } + + #[test] + fn overlay_embedding_prefixes_some_wins_over_a_lower_layer() { + let base = Figment::from(Serialized::defaults(Config::default())).merge( + Serialized::default("embedding_query_prefix", "toml-query: "), + ); + let overlaid = overlay_embedding_prefixes(base, Some("env-query: "), None); + let cfg: Config = overlaid.extract().unwrap(); + assert_eq!(cfg.embedding_query_prefix.as_deref(), Some("env-query: ")); + } + + #[test] + fn overlay_embedding_prefixes_empty_string_clears_a_lower_layer() { + // Present but empty (`Some("")`) is a deliberate override — an + // operator clearing a `config.toml` value via env without editing + // the file — distinct from `None` (the previous test), which must + // leave the lower layer untouched. + let base = Figment::from(Serialized::defaults(Config::default())).merge( + Serialized::default("embedding_query_prefix", "toml-query: "), + ); + let overlaid = overlay_embedding_prefixes(base, Some(""), None); + let cfg: Config = overlaid.extract().unwrap(); + assert_eq!(cfg.embedding_query_prefix.as_deref(), Some("")); + } + + /// Exercises the real `Config::load` end to end, through an absolute + /// `config.toml` path and data dir from a `TempDir`. TOML strings were + /// never subject to figment's `Env`-provider trimming in the first + /// place, so this path already worked before the fix; this guards it + /// staying correct. + /// + /// This process's own environment is shared with every other test in + /// this binary and could already carry one of the two prefix vars from + /// the test runner's shell, which would make `Config::load` pick the + /// env value over the TOML one below and this test would silently stop + /// verifying the TOML-only path. Rather than assume the ambient + /// environment is clean, the actual `Config::load` call runs in a + /// separate child process with both vars explicitly removed via + /// `Command::env_remove` — real isolation instead of an in-process + /// assumption, and it does not touch this rule's target (mutating + /// *this* process's env), since a spawned child's environment is its + /// own. + #[test] + fn loader_toml_path_preserves_whitespace_with_no_env_var_set() { + const CHILD_MARKER: &str = "AI_MEMORY_TEST_LOADER_TOML_PATH_CHILD"; + if std::env::var_os(CHILD_MARKER).is_some() { + // Running as the child, with both prefix vars removed by the + // parent below: do the real work and print the result for the + // parent to assert on. + let tmp = TempDir::new().unwrap(); + let config_path = tmp.path().join("config.toml"); + std::fs::write( + &config_path, + "embedding_query_prefix = \"query: \"\n\ + embedding_document_prefix = \"passage: \"\n", + ) + .unwrap(); + let cfg = Config::load(Some(&config_path), Some(tmp.path().to_path_buf())).unwrap(); + println!( + "query={:?} document={:?}", + cfg.embedding_query_prefix, cfg.embedding_document_prefix + ); + return; + } + // Running as the parent: re-exec this same test binary, filtered + // to just this one test, as a genuinely separate child process + // with both prefix env vars removed. + let exe = std::env::current_exe().expect("current test binary path"); + let output = std::process::Command::new(&exe) + .arg("--exact") + .arg("config::tests::loader_toml_path_preserves_whitespace_with_no_env_var_set") + .arg("--test-threads=1") + .arg("--nocapture") + .env(CHILD_MARKER, "1") + .env_remove("AI_MEMORY_EMBEDDING_QUERY_PREFIX") + .env_remove("AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX") + .output() + .expect("failed to spawn child test process"); + assert!( + output.status.success(), + "child test process failed:\nstdout: {}\nstderr: {}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + stdout.contains("query=Some(\"query: \")"), + "child stdout: {stdout}" + ); + assert!( + stdout.contains("document=Some(\"passage: \")"), + "child stdout: {stdout}" + ); + } + #[test] fn copilot_embedding_defaults_model_dim_and_reuses_copilot_auth() { let tmp = TempDir::new().unwrap(); diff --git a/crates/ai-memory-cli/tests/suite/packaging.rs b/crates/ai-memory-cli/tests/suite/packaging.rs index db61b258e..a2ab3675c 100644 --- a/crates/ai-memory-cli/tests/suite/packaging.rs +++ b/crates/ai-memory-cli/tests/suite/packaging.rs @@ -1014,7 +1014,16 @@ fn run_wrapper_with_fake_docker_env( .env("AI_MEMORY_DATA_VOLUME", "test-ai-memory-data") .env("HOME", shell_path(tmp.path())) .env_remove("AI_MEMORY_SERVER_URL") - .env_remove("CLAUDE_CONFIG_DIR"); + .env_remove("CLAUDE_CONFIG_DIR") + // Guarantees the "unset" case in + // `posix_wrapper_forwards_embedding_prefixes_by_presence_not_non_emptiness` + // is actually unset rather than silently inheriting whatever the + // test-runner's own ambient environment happens to hold; every + // other case re-adds one of these via `forwarded_env` below, so + // removing them unconditionally here is a no-op for every other + // caller of this helper. + .env_remove("AI_MEMORY_EMBEDDING_QUERY_PREFIX") + .env_remove("AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX"); if let Some(claude_config_dir) = claude_config_dir { command.env("CLAUDE_CONFIG_DIR", claude_config_dir); } @@ -1833,6 +1842,83 @@ mod slow { } } + /// The two embedding-prefix env vars are forwarded on PRESENCE, not + /// non-emptiness, unlike every other var in the loop: an operator sets + /// one to the empty string to clear a `config.toml`-configured prefix + /// without editing the file (see `Config::load`'s figment overlay in + /// `ai-memory-cli/src/config.rs`), and that override only reaches the + /// server if the wrapper forwards the (empty) variable rather than + /// dropping it the way a plain `[ -n ]` check would. + #[cfg(unix)] + #[test] + fn posix_wrapper_forwards_embedding_prefixes_by_presence_not_non_emptiness() { + const QUERY: &str = "AI_MEMORY_EMBEDDING_QUERY_PREFIX"; + const DOC: &str = "AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX"; + let has_e = |args: &[&str], name: &str| args.windows(2).any(|pair| pair == ["-e", name]); + + // Unset: neither var forwarded (the loop must not invent a value). + let args = + run_wrapper_with_fake_docker_and_forwarded_env(&["llm-test"], "[name=seccomp]", &[]); + let lines: Vec<&str> = args.lines().collect(); + assert!( + !has_e(&lines, QUERY), + "unset must not be forwarded; got {lines:?}" + ); + assert!( + !has_e(&lines, DOC), + "unset must not be forwarded; got {lines:?}" + ); + + // Empty: forwarded anyway — this is the override case. + let args = run_wrapper_with_fake_docker_and_forwarded_env( + &["llm-test"], + "[name=seccomp]", + &[(QUERY, ""), (DOC, "")], + ); + let lines: Vec<&str> = args.lines().collect(); + assert!( + has_e(&lines, QUERY), + "an empty (but present) value must still be forwarded; got {lines:?}" + ); + assert!( + has_e(&lines, DOC), + "an empty (but present) value must still be forwarded; got {lines:?}" + ); + + // Whitespace-only: also present, also forwarded — this loop must + // not apply any trimming/emptiness judgement of its own. + let args = run_wrapper_with_fake_docker_and_forwarded_env( + &["llm-test"], + "[name=seccomp]", + &[(QUERY, " "), (DOC, " ")], + ); + let lines: Vec<&str> = args.lines().collect(); + assert!( + has_e(&lines, QUERY), + "whitespace-only must still be forwarded; got {lines:?}" + ); + assert!( + has_e(&lines, DOC), + "whitespace-only must still be forwarded; got {lines:?}" + ); + + // Non-empty: forwarded, same as every other var. + let args = run_wrapper_with_fake_docker_and_forwarded_env( + &["llm-test"], + "[name=seccomp]", + &[(QUERY, "query: "), (DOC, "passage: ")], + ); + let lines: Vec<&str> = args.lines().collect(); + assert!( + has_e(&lines, QUERY), + "a non-empty value must be forwarded; got {lines:?}" + ); + assert!( + has_e(&lines, DOC), + "a non-empty value must be forwarded; got {lines:?}" + ); + } + // The Windows mirror of macos_wrapper_routes_urls_by_real_subcommand: Docker // Desktop gives Linux containers no host networking on Windows either, so the // helper container cannot reach the host-published server over loopback. diff --git a/crates/ai-memory-consolidate/src/embed.rs b/crates/ai-memory-consolidate/src/embed.rs index 57299d2c3..0302fac6f 100644 --- a/crates/ai-memory-consolidate/src/embed.rs +++ b/crates/ai-memory-consolidate/src/embed.rs @@ -91,7 +91,10 @@ pub async fn run_embedding_backfill( options: EmbedBackfillOptions, ) -> Result { let provider = embedder.provider().to_string(); - let model = embedder.model().to_string(); + // Not `.model()`: this identity drives both stale-row detection below + // (a document-prefix change must look "stale" like a model change + // would) and every write this pass makes. See `Embedder::model_identity`. + let model = embedder.model_identity(); let dim = embedder.dim(); let candidates = reader.decay_candidates(workspace_id, project_id).await?; diff --git a/crates/ai-memory-consolidate/tests/suite/embed_backfill.rs b/crates/ai-memory-consolidate/tests/suite/embed_backfill.rs index e8e13cf84..3864e4d8f 100644 --- a/crates/ai-memory-consolidate/tests/suite/embed_backfill.rs +++ b/crates/ai-memory-consolidate/tests/suite/embed_backfill.rs @@ -156,3 +156,272 @@ async fn backfill_reembed_regenerates_existing_rows() { assert_eq!(counts.skipped, 0); assert_eq!(f.embedded_count().await, 1); } + +/// A fake embedder standing in for `OpenAiEmbedder`/`OpenAiCompatEmbedder` +/// with a configurable document prefix: `identity_suffix` plays the role a +/// non-empty `document_prefix` plays there (`Embedder::model_identity` +/// diverges from `Embedder::model` when set), and `fail_on_substring` lets +/// one test drive a genuine per-page provider failure without a real +/// network call. +struct ConfigurableEmbedder { + dim: u32, + identity_suffix: String, + fail_on_substring: Option<&'static str>, +} + +#[async_trait::async_trait] +impl Embedder for ConfigurableEmbedder { + fn provider(&self) -> &'static str { + "synthetic" + } + + fn model(&self) -> &str { + "bag-of-words-v1" + } + + fn model_identity(&self) -> String { + if self.identity_suffix.is_empty() { + self.model().to_string() + } else { + format!("{}+{}", self.model(), self.identity_suffix) + } + } + + fn dim(&self) -> u32 { + self.dim + } + + async fn embed(&self, text: &str) -> ai_memory_llm::LlmResult> { + if let Some(marker) = self.fail_on_substring + && text.contains(marker) + { + return Err(ai_memory_llm::LlmError::Provider { + status: 500, + body: "synthetic configured failure".into(), + }); + } + Ok(vec![0.5_f32; self.dim as usize]) + } +} + +/// A document-prefix change (a new `identity_suffix` here) must make a +/// **normal** (non-force) backfill pass treat every page as needing a +/// fresh embedding, exactly like a model change would — this is what +/// keeps `memory_query`'s retrieval eligibility filter (keyed on the +/// currently configured embedder's `model_identity`) from ever mixing +/// vectors embedded under different prefixes. The old row is left in +/// place (nothing here deletes it — deletion is a separate, explicit step +/// mirroring `POST /admin/embed`'s `reembed` purge, covered by the next +/// test), but it is no longer reachable under the newly configured +/// identity. +#[tokio::test] +async fn document_prefix_change_makes_backfill_treat_pages_as_stale_and_reembeds_them() { + let f = fixture().await; + f.write_page("notes/a.md", "alpha content").await; + + let legacy: Arc = Arc::new(ConfigurableEmbedder { + dim: 8, + identity_suffix: String::new(), + fail_on_substring: None, + }); + let first = run_embedding_backfill( + &f.store.reader, + &f.store.writer, + &f.wiki, + &legacy, + f.ws, + f.proj, + EmbedBackfillOptions::default(), + ) + .await + .expect("first backfill"); + assert_eq!(first.embedded, 1); + + // Simulates reconfiguring `embedding_document_prefix`: same provider, + // same wire model, same dim — only the document-side identity differs. + let prefixed: Arc = Arc::new(ConfigurableEmbedder { + dim: 8, + identity_suffix: "dp-test".into(), + fail_on_substring: None, + }); + assert_ne!(legacy.model_identity(), prefixed.model_identity()); + + // Between the config change and the next backfill pass, the stored + // row is still under the OLD identity — this is the property the fix + // protects: a query using the NOW-configured embedder's identity (what + // `memory_query`'s retrieval eligibility filter does) does not match + // it, so the page is excluded from vector search rather than silently + // matched with content embedded under the old prefix. + let prefixed_ids_before = f + .store + .reader + .embedded_page_ids( + f.ws, + f.proj, + prefixed.provider().to_string(), + prefixed.model_identity(), + prefixed.dim(), + ) + .await + .expect("prefixed embedded_page_ids before backfill"); + assert_eq!( + prefixed_ids_before.len(), + 0, + "the pre-change row must not be matched under the new identity" + ); + + let second = run_embedding_backfill( + &f.store.reader, + &f.store.writer, + &f.wiki, + &prefixed, + f.ws, + f.proj, + EmbedBackfillOptions::default(), // reembed: false — a NORMAL pass + ) + .await + .expect("backfill after a document-prefix change"); + assert_eq!( + second.embedded, 1, + "a document-prefix change must be treated as stale and re-embedded, \ + not silently skipped as \"already embedded\"" + ); + assert_eq!(second.skipped, 0); + + // `page_embeddings` is one row per page (replace-on-write), so the + // fresh write replaces the pre-change row rather than accumulating + // alongside it: the old identity now matches nothing … + let legacy_ids_after = f + .store + .reader + .embedded_page_ids( + f.ws, + f.proj, + legacy.provider().to_string(), + legacy.model_identity(), + legacy.dim(), + ) + .await + .expect("legacy embedded_page_ids after backfill"); + assert_eq!( + legacy_ids_after.len(), + 0, + "the pre-change row was replaced, not left behind, by the re-embed" + ); + // … and the new identity now does. + let prefixed_ids_after = f + .store + .reader + .embedded_page_ids( + f.ws, + f.proj, + prefixed.provider().to_string(), + prefixed.model_identity(), + prefixed.dim(), + ) + .await + .expect("prefixed embedded_page_ids after backfill"); + assert_eq!(prefixed_ids_after.len(), 1); +} + +/// Mirrors `POST /admin/embed`'s `reembed: true` sequence exactly +/// (`ai-memory-mcp/src/admin.rs::handle_embed`): purge stale rows for the +/// configured triple, then run a forced backfill. One of two pages fails +/// mid-run; the purge must still have applied to both, the surviving page +/// must still get a fresh row, and the failed page must be left with NO +/// row (not a stale one and not a half-written one) rather than silently +/// keeping a purged-then-never-replaced gap invisible to the operator — +/// `counts.failed` is what surfaces it. +#[tokio::test] +async fn force_rebuild_purges_first_and_counts_a_partial_failure_correctly() { + let f = fixture().await; + f.write_page("notes/ok.md", "alpha content").await; + f.write_page("notes/bad.md", "FAIL_MARKER content").await; + + // Seed both pages under a stand-in "old" identity, matching a real + // force-rebuild's starting point: pages already have rows, possibly + // under a different identity than the one about to be forced. + let legacy: Arc = Arc::new(ConfigurableEmbedder { + dim: 8, + identity_suffix: "old".into(), + fail_on_substring: None, + }); + let seeded = run_embedding_backfill( + &f.store.reader, + &f.store.writer, + &f.wiki, + &legacy, + f.ws, + f.proj, + EmbedBackfillOptions::default(), + ) + .await + .expect("seed backfill"); + assert_eq!(seeded.embedded, 2); + + let flaky: Arc = Arc::new(ConfigurableEmbedder { + dim: 8, + identity_suffix: String::new(), // the configured (new) identity + fail_on_substring: Some("FAIL_MARKER"), + }); + + // Step 1 (admin.rs's purge-before-reembed): delete every row in scope + // that does not match the target triple — here that is BOTH seeded + // rows, since they are under `legacy`'s identity. + let purged = f + .store + .writer + .delete_stale_page_embeddings( + f.ws, + Some(f.proj), + flaky.provider().to_string(), + flaky.model_identity(), + flaky.dim(), + ) + .await + .expect("purge"); + assert_eq!( + purged, 2, + "purge removes every row under a non-matching identity" + ); + + // Step 2: forced backfill with the embedder that fails on notes/bad.md. + let counts = run_embedding_backfill( + &f.store.reader, + &f.store.writer, + &f.wiki, + &flaky, + f.ws, + f.proj, + EmbedBackfillOptions { + reembed: true, + dry_run: false, + }, + ) + .await + .expect("forced backfill with a partial failure"); + assert_eq!(counts.embedded, 1, "notes/ok.md succeeds"); + assert_eq!(counts.failed, 1, "notes/bad.md's provider call fails"); + + let current_ids = f + .store + .reader + .embedded_page_ids( + f.ws, + f.proj, + flaky.provider().to_string(), + flaky.model_identity(), + flaky.dim(), + ) + .await + .expect("embedded_page_ids after the forced rebuild"); + assert_eq!( + current_ids.len(), + 1, + "exactly the succeeding page has a row under the new identity; \ + the failed page is left with none — purged, not silently kept \ + stale, and not half-written — so it surfaces as `failed` for the \ + operator to retry rather than looking either fully migrated or \ + still on the old prefix" + ); +} diff --git a/crates/ai-memory-llm/src/embedding.rs b/crates/ai-memory-llm/src/embedding.rs index f2abc1ba5..52cf2a211 100644 --- a/crates/ai-memory-llm/src/embedding.rs +++ b/crates/ai-memory-llm/src/embedding.rs @@ -18,6 +18,7 @@ use std::time::Duration; use async_trait::async_trait; use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use tracing::debug; use crate::error::{LlmError, LlmResult}; @@ -39,9 +40,35 @@ pub trait Embedder: Send + Sync { /// Short identifier (e.g. `openai`, `voyage`, `synthetic`). fn provider(&self) -> &'static str; - /// Model identifier (e.g. `text-embedding-3-small`). + /// Model identifier (e.g. `text-embedding-3-small`). The exact string + /// sent on the wire — never suffixed or altered. See + /// [`Self::model_identity`] for the value that should be stored + /// alongside a vector or used to select eligible stored vectors. fn model(&self) -> &str; + /// The `model` component of the stored/matched `(provider, model, dim)` + /// identity for a **document** embedding — distinct from [`Self::model`] + /// (the wire model name) whenever the effective document-side input + /// text is not what the model id alone implies. Defaults to + /// `self.model().to_string()`: providers with no such distinction (every + /// provider except the two below) are unaffected. + /// + /// `OpenAiEmbedder`/`OpenAiCompatEmbedder` override this to fold in a + /// fingerprint of their configured `document_prefix` when it is + /// non-empty, so a page embedded under one document prefix is never + /// silently matched against, or mixed with, vectors embedded under a + /// different (or no) prefix — the same `(provider, model, dim)` + /// refuse-on-mismatch and stale-row machinery that already protects + /// against a plain model swap now also protects against a document + /// prefix change. An empty prefix reproduces `self.model()` exactly + /// (the pre-existing identity), so upgrading installs with no prefix + /// configured need no migration. The **query** prefix never affects + /// this: only the text actually embedded and stored needs a distinct + /// identity, and a query is never stored. + fn model_identity(&self) -> String { + self.model().to_string() + } + /// Vector dimensionality. fn dim(&self) -> u32; @@ -67,6 +94,8 @@ pub struct OpenAiEmbedder { base_url: String, model: String, dim: u32, + query_prefix: String, + document_prefix: String, } impl OpenAiEmbedder { @@ -90,6 +119,8 @@ impl OpenAiEmbedder { base_url: "https://api.openai.com".into(), model: model.into(), dim, + query_prefix: String::new(), + document_prefix: String::new(), }) } @@ -99,6 +130,20 @@ impl OpenAiEmbedder { self.base_url = url.into(); self } + + /// Set the query/document prefixes prepended before embedding (see + /// [`OpenAiCompatEmbedder::with_prefixes`]). Empty strings are a no-op, + /// so this is safe to call unconditionally with the configured values. + #[must_use] + pub fn with_prefixes( + mut self, + query_prefix: impl Into, + document_prefix: impl Into, + ) -> Self { + self.query_prefix = query_prefix.into(); + self.document_prefix = document_prefix.into(); + self + } } #[derive(Debug, Serialize)] @@ -262,6 +307,10 @@ impl Embedder for OpenAiEmbedder { &self.model } + fn model_identity(&self) -> String { + document_prefix_identity(&self.model, &self.document_prefix) + } + fn dim(&self) -> u32 { self.dim } @@ -277,6 +326,32 @@ impl Embedder for OpenAiEmbedder { ) .await } + + async fn embed_document(&self, text: &str) -> LlmResult> { + let prefixed = prepend_prefix(&self.document_prefix, text); + openai_style_embed( + &self.client, + &self.base_url, + Some(&self.api_key), + &self.model, + self.dim, + &prefixed, + ) + .await + } + + async fn embed_query(&self, text: &str) -> LlmResult> { + let prefixed = prepend_prefix(&self.query_prefix, text); + openai_style_embed( + &self.client, + &self.base_url, + Some(&self.api_key), + &self.model, + self.dim, + &prefixed, + ) + .await + } } /// OpenAI-compatible embeddings endpoint (Ollama / LM Studio / vLLM / @@ -291,6 +366,21 @@ pub struct OpenAiCompatEmbedder { base_url: String, model: String, dim: u32, + /// Prepended to every query text before embedding (before truncation). + /// Empty by default: symmetric models (most OpenAI-compatible servers) + /// see no behaviour change. Asymmetric models need a query-side + /// instruction the OpenAI-compatible `/v1/embeddings` wire format has + /// no field for — the client has to prepend it instead. + /// `nvidia/Nemotron-3-Embed-1B-BF16` and base E5 models use a simple + /// `"query: "` / `"passage: "` pair; `e5-mistral-7b-instruct` and + /// Qwen3-Embedding instead need a full task-instruction string with + /// different exact spacing each (documents plain for both — see + /// `ai-memory-cli/src/config.rs`'s `embedding_query_prefix` field doc + /// comment for the two exact templates). See `with_prefixes`. + query_prefix: String, + /// Document-side counterpart of `query_prefix` (e.g. `"passage: "` for + /// Nemotron-3-Embed / base E5 — not every model needs one). + document_prefix: String, } impl OpenAiCompatEmbedder { @@ -316,8 +406,30 @@ impl OpenAiCompatEmbedder { base_url: base_url.into(), model: model.into(), dim, + query_prefix: String::new(), + document_prefix: String::new(), }) } + + /// Set the strings prepended to query / document text before embedding, + /// applied ahead of the existing truncation so truncation still bounds + /// the whole request body (prefix included). Pass empty strings for no + /// prefix (the default); safe to call unconditionally with a possibly + /// unset operator setting. + /// + /// Publisher-specified prefixes are exact strings, often with a + /// significant trailing space (e.g. Nemotron-3-Embed / the E5 family + /// use `"query: "` and `"passage: "`) — callers must not trim them. + #[must_use] + pub fn with_prefixes( + mut self, + query_prefix: impl Into, + document_prefix: impl Into, + ) -> Self { + self.query_prefix = query_prefix.into(); + self.document_prefix = document_prefix.into(); + self + } } #[async_trait] @@ -330,6 +442,10 @@ impl Embedder for OpenAiCompatEmbedder { &self.model } + fn model_identity(&self) -> String { + document_prefix_identity(&self.model, &self.document_prefix) + } + fn dim(&self) -> u32 { self.dim } @@ -345,6 +461,32 @@ impl Embedder for OpenAiCompatEmbedder { ) .await } + + async fn embed_document(&self, text: &str) -> LlmResult> { + let prefixed = prepend_prefix(&self.document_prefix, text); + openai_style_embed( + &self.client, + &self.base_url, + self.api_key.as_ref(), + &self.model, + self.dim, + &prefixed, + ) + .await + } + + async fn embed_query(&self, text: &str) -> LlmResult> { + let prefixed = prepend_prefix(&self.query_prefix, text); + openai_style_embed( + &self.client, + &self.base_url, + self.api_key.as_ref(), + &self.model, + self.dim, + &prefixed, + ) + .await + } } /// Voyage Embeddings API. @@ -501,6 +643,49 @@ impl Embedder for SyntheticEmbedder { } } +/// Prepend `prefix` to `text` for an asymmetric embedding model, applied +/// BEFORE [`truncate_for_embedding`] so truncation always bounds the whole +/// request body (prefix included) rather than letting a prefix push the +/// text itself past the server's token limit. An empty prefix — the +/// default — returns `text` unchanged and borrowed, so the unset case +/// allocates nothing. +pub(crate) fn prepend_prefix<'a>(prefix: &str, text: &'a str) -> std::borrow::Cow<'a, str> { + if prefix.is_empty() { + std::borrow::Cow::Borrowed(text) + } else { + std::borrow::Cow::Owned(format!("{prefix}{text}")) + } +} + +/// The `model` component of the stored embedding identity for a document +/// embedder configured with `document_prefix`. An empty prefix returns +/// `model` unchanged — the pre-existing identity, so an install with no +/// document prefix configured needs no migration. A non-empty prefix +/// appends a versioned SHA-256 fingerprint of the prefix bytes: `+dp1-` +/// followed by the first 16 hex characters (64 bits) of +/// `SHA-256(document_prefix)`. Not the raw prefix text itself, which +/// could be long, contain characters awkward in a stored column, or leak +/// an operator's exact instruction string into logs/admin output more +/// than necessary — and not a 32-bit hash (an earlier revision used +/// `fnv1a` truncated to 32 bits, which a search for a same-length ASCII +/// collision found in minutes: `"Document category 5hvhw0: "` and +/// `"Document category 1i4yh8i: "` both fingerprinted to `5c7f37a2`, +/// which would have silently mixed two different prefixes' vectors under +/// one identity — see `document_prefix_identity_does_not_collide_on_the_known_fnv32_pair` +/// below). SHA-256 truncated to 64 bits keeps the collision probability +/// negligible for the small number of prefixes one deployment actually +/// configures over time, without needing the full 256-bit digest in a +/// column meant to stay human-scannable. The `dp1` version tag lets a +/// future change to this scheme be distinguished from today's rather than +/// risking a silent collision with an old identity under a new one. +pub(crate) fn document_prefix_identity(model: &str, document_prefix: &str) -> String { + if document_prefix.is_empty() { + return model.to_string(); + } + let digest = format!("{:x}", Sha256::digest(document_prefix.as_bytes())); + format!("{model}+dp1-{}", &digest[..16]) +} + /// Unit-normalise so dot-product equals cosine similarity. pub(crate) fn normalise(mut v: Vec) -> Vec { let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); @@ -532,6 +717,144 @@ pub fn cosine(a: &[f32], b: &[f32]) -> f32 { mod tests { use super::*; + #[test] + fn document_prefix_identity_is_unchanged_when_the_prefix_is_empty() { + // The legacy-identity guarantee: an install with no document + // prefix configured must key its vectors exactly as before this + // feature, so it needs no migration. + assert_eq!( + document_prefix_identity("nomic-embed-text", ""), + "nomic-embed-text" + ); + } + + #[test] + fn document_prefix_identity_changes_deterministically_with_the_prefix() { + let base = document_prefix_identity("nomic-embed-text", "passage: "); + assert_ne!( + base, "nomic-embed-text", + "a set prefix must not reuse the legacy identity" + ); + assert!(base.starts_with("nomic-embed-text+dp1-")); + // Deterministic: the same (model, prefix) always produces the same + // identity, so pages embedded in different requests still land + // under one queryable identity. + assert_eq!( + base, + document_prefix_identity("nomic-embed-text", "passage: ") + ); + } + + #[test] + fn document_prefix_identity_distinguishes_different_prefixes() { + // Two distinct document prefixes on the same model must not + // collide — otherwise a prefix *change* would look like no change + // at all to the refuse-on-mismatch / stale-row machinery. + let a = document_prefix_identity("nomic-embed-text", "passage: "); + let b = document_prefix_identity("nomic-embed-text", "document: "); + assert_ne!(a, b); + } + + /// Regression test for a real collision found in an earlier revision + /// of `document_prefix_identity`, which truncated `fnv1a` to 32 bits: + /// `"Document category 5hvhw0: "` and `"Document category 1i4yh8i: "` + /// both fingerprinted to `5c7f37a2` (verified independently in + /// Python), which would have silently mixed the two prefixes' vectors + /// under one stored identity. The SHA-256-based fingerprint here does + /// not collide on this pair (also independently verified: their first + /// 16 hex characters are `b32006f5a41f3d24` and `513594f86d5883fd`). + #[test] + fn document_prefix_identity_does_not_collide_on_the_known_fnv32_pair() { + let a = document_prefix_identity("nomic-embed-text", "Document category 5hvhw0: "); + let b = document_prefix_identity("nomic-embed-text", "Document category 1i4yh8i: "); + assert_ne!( + a, b, + "these two prefixes collided under the old 32-bit fnv1a fingerprint; \ + the new SHA-256-based one must not repeat that collision" + ); + } + + #[test] + fn model_identity_defaults_to_model_for_providers_without_prefixes() { + // Providers that never got the document-prefix override (google, + // voyage, local, copilot, synthetic) must see byte-identical + // behaviour: `model_identity` defaults to `model().to_string()`. + let e = SyntheticEmbedder::new(8); + assert_eq!(e.model_identity(), e.model()); + } + + #[test] + fn openai_compat_embedder_model_identity_reflects_the_document_prefix() { + let unset = OpenAiCompatEmbedder::new("http://localhost:9/v1", None, "nomic-embed-text", 8) + .expect("embedder builds"); + assert_eq!(unset.model_identity(), "nomic-embed-text"); + + let with_prefix = + OpenAiCompatEmbedder::new("http://localhost:9/v1", None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + assert_ne!(with_prefix.model_identity(), "nomic-embed-text"); + assert_eq!( + with_prefix.model_identity(), + document_prefix_identity("nomic-embed-text", "passage: ") + ); + // The query prefix must NOT affect the stored identity — only + // documents are stored; a query is never persisted. + let query_only = + OpenAiCompatEmbedder::new("http://localhost:9/v1", None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", ""); + assert_eq!(query_only.model_identity(), "nomic-embed-text"); + } + + #[test] + fn openai_embedder_model_identity_reflects_the_document_prefix() { + let unset = + OpenAiEmbedder::new(SecretString::from("k"), "text-embedding-3-small", 1536).unwrap(); + assert_eq!(unset.model_identity(), "text-embedding-3-small"); + + let with_prefix = + OpenAiEmbedder::new(SecretString::from("k"), "text-embedding-3-small", 1536) + .unwrap() + .with_prefixes("query: ", "passage: "); + assert_ne!(with_prefix.model_identity(), "text-embedding-3-small"); + } + + /// Transport-level proof that `OpenAiEmbedder` (not just + /// `OpenAiCompatEmbedder`, covered in + /// `tests/suite/openai_compat_embedder.rs`) actually sends the + /// configured prefix on the wire — `with_prefixes` alone only proves + /// the fields are stored, not that `embed_query`/`embed_document` use + /// them in the real HTTP request body. + #[tokio::test] + async fn openai_embedder_sends_the_prefix_on_the_wire() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap(); + assert_eq!(body["input"], "query: find the runbook"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "data": [{ "embedding": vec![0.5_f32; 4] }], + })) + }) + .expect(1) + .mount(&server) + .await; + + let e = OpenAiEmbedder::new(SecretString::from("sk-test"), "text-embedding-3-small", 4) + .unwrap() + .with_base_url(server.uri()) + .with_prefixes("query: ", "passage: "); + + e.embed_query("find the runbook") + .await + .expect("embed_query succeeds"); + } + #[tokio::test] async fn synthetic_embedder_produces_unit_vectors() { let e = SyntheticEmbedder::new(64); @@ -592,4 +915,65 @@ mod tests { let err = parse_openai_embedding_values(body, 200).unwrap_err(); assert!(matches!(err, LlmError::Provider { status: 200, .. })); } + + #[test] + fn prepend_prefix_applies_publisher_instruction() { + // Nemotron-3-Embed / the E5 family: exact publisher strings, + // including the significant trailing space. + assert_eq!( + prepend_prefix("query: ", "find the release runbook"), + "query: find the release runbook" + ); + assert_eq!( + prepend_prefix("passage: ", "the release runbook says..."), + "passage: the release runbook says..." + ); + } + + #[test] + fn prepend_prefix_unset_leaves_text_unchanged() { + // The default (empty prefix) must be a true no-op, byte for byte, + // and borrow rather than allocate. + let text = "unchanged text"; + let out = prepend_prefix("", text); + assert_eq!(out, text); + assert!(matches!(out, std::borrow::Cow::Borrowed(_))); + } + + #[test] + fn prefix_is_applied_before_truncation_so_it_still_bounds_the_whole_input() { + // A prefix pushes the *total* input closer to the cap; truncation + // must run on prefix+text together, never on text alone with the + // prefix appended afterwards (which could exceed the server limit). + let long_text = "x".repeat(50_000); + let prefixed = prepend_prefix("passage: ", &long_text); + let truncated = truncate_for_embedding(&prefixed, OPENAI_EMBED_MAX_TOKENS); + assert!(truncated.starts_with("passage: ")); + assert!(truncated.len() <= 8_000, "must respect the hard byte cap"); + assert!(truncated.ends_with('…')); + } + + #[test] + fn openai_compat_embedder_defaults_to_no_prefix() { + // Construction without `with_prefixes` must be byte-identical to + // before this feature existed. + let e = OpenAiCompatEmbedder::new("http://localhost:9/v1", None, "nomic-embed-text", 8) + .expect("embedder builds"); + assert_eq!(e.query_prefix, ""); + assert_eq!(e.document_prefix, ""); + } + + #[test] + fn openai_compat_embedder_stores_configured_prefixes() { + let e = OpenAiCompatEmbedder::new( + "http://localhost:9/v1", + None, + "nvidia/Nemotron-3-Embed-1B-BF16", + 2048, + ) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + assert_eq!(e.query_prefix, "query: "); + assert_eq!(e.document_prefix, "passage: "); + } } diff --git a/crates/ai-memory-llm/src/factory.rs b/crates/ai-memory-llm/src/factory.rs index 7ad56ddc6..a5d6b208a 100644 --- a/crates/ai-memory-llm/src/factory.rs +++ b/crates/ai-memory-llm/src/factory.rs @@ -206,6 +206,23 @@ pub struct EmbedderConfig { /// a warning instead of refusing to start; an explicitly configured /// one still fails hard. pub defaulted: bool, + /// Prepended to every query text before embedding, ahead of the + /// existing truncation. Empty (the default) is a no-op. Only the + /// `openai` and `openai-compat` embedders apply it; other providers + /// (`google` has its own built-in task-type asymmetry, `voyage`, + /// `local`, `copilot`) ignore it. Needed for asymmetric self-hosted + /// models whose publisher specifies a query-side instruction the + /// OpenAI-compatible `/v1/embeddings` wire format has no field for. + /// `nvidia/Nemotron-3-Embed-1B-BF16` and base E5 models + /// (`intfloat/e5-base-v2`, multilingual E5, …) use a simple + /// `"query: "` string; instruction-tuned E5 variants and + /// Qwen3-Embedding instead need a full task-instruction string (their + /// documents stay plain — leave `document_prefix` unset for those). + pub query_prefix: String, + /// Document-side counterpart of `query_prefix` (e.g. `"passage: "` for + /// Nemotron-3-Embed / base E5 — not every model needs one; see + /// `query_prefix`'s doc comment). + pub document_prefix: String, } /// Construct an `Arc` from the config. @@ -221,7 +238,8 @@ pub fn build_embedder(config: EmbedderConfig) -> LlmResult> { } let arc: Arc = match config.provider { EmbedderChoice::OpenAi => { - let mut e = OpenAiEmbedder::new(config.api_key, config.model, config.dim)?; + let mut e = OpenAiEmbedder::new(config.api_key, config.model, config.dim)? + .with_prefixes(config.query_prefix, config.document_prefix); if let Some(url) = config.base_url { e = e.with_base_url(url); } @@ -246,12 +264,10 @@ pub fn build_embedder(config: EmbedderConfig) -> LlmResult> { .base_url .ok_or_else(|| LlmError::NotConfigured("AI_MEMORY_EMBEDDING_BASE_URL".into()))?; let api_key = (!config.api_key.expose_secret().is_empty()).then_some(config.api_key); - Arc::new(OpenAiCompatEmbedder::new( - base, - api_key, - config.model, - config.dim, - )?) + Arc::new( + OpenAiCompatEmbedder::new(base, api_key, config.model, config.dim)? + .with_prefixes(config.query_prefix, config.document_prefix), + ) } #[cfg(feature = "local-embeddings")] EmbedderChoice::Local => { diff --git a/crates/ai-memory-llm/src/health.rs b/crates/ai-memory-llm/src/health.rs index 6fd7ceb39..ceb530405 100644 --- a/crates/ai-memory-llm/src/health.rs +++ b/crates/ai-memory-llm/src/health.rs @@ -375,6 +375,17 @@ impl Embedder for HealthRecordingEmbedder { self.inner.model() } + fn model_identity(&self) -> String { + // Not the default (`self.model().to_string()`): every configured + // embedder is wrapped in this type before being handed to the rest + // of the server (`serve.rs`'s `wrap_embedder`), so without this + // override every caller of `model_identity` — the refuse-on-mismatch + // check, backfill, retrieval, cleanup — would silently see the + // wire model name instead of the inner embedder's actual document- + // prefix-aware identity. + self.inner.model_identity() + } + fn dim(&self) -> u32 { self.inner.dim() } @@ -559,6 +570,35 @@ mod tests { assert!(after.last_call_at.is_some()); } + /// Every configured embedder is wrapped in `HealthRecordingEmbedder` + /// before reaching the rest of the server (`serve.rs`'s + /// `wrap_embedder`), so a caller of `model_identity` — the + /// refuse-on-mismatch check, backfill, retrieval, cleanup — only ever + /// sees the wrapper, never the inner embedder directly. Without the + /// wrapper's own `model_identity` override, the default trait impl + /// (`self.model().to_string()`) would silently discard a document + /// prefix's identity fingerprint. `TaskAwareEmbedder` doesn't + /// distinguish this (it never overrides `model_identity` either, so a + /// missing wrapper override would coincidentally still pass against + /// it); a real `OpenAiCompatEmbedder` with a document prefix set does. + #[test] + fn embedder_wrapper_forwards_the_inner_model_identity_override() { + let inner = + crate::OpenAiCompatEmbedder::new("http://localhost:9/v1", None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + let inner_identity = inner.model_identity(); + assert_ne!( + inner_identity, + inner.model(), + "the fixture must actually have a distinct identity, or this test proves nothing" + ); + + let health = ProviderHealth::default(); + let wrapped = health.wrap_embedder(Arc::new(inner), "openai-compat", "nomic-embed-text", 8); + assert_eq!(wrapped.model_identity(), inner_identity); + } + #[test] fn provider_health_status_wire_labels() { for (status, label) in [ diff --git a/crates/ai-memory-llm/tests/suite/openai_compat_embedder.rs b/crates/ai-memory-llm/tests/suite/openai_compat_embedder.rs index 2cd5021be..08f2e3708 100644 --- a/crates/ai-memory-llm/tests/suite/openai_compat_embedder.rs +++ b/crates/ai-memory-llm/tests/suite/openai_compat_embedder.rs @@ -87,6 +87,8 @@ async fn factory_builds_compat_embedder_and_requires_base_url() { models_dir: None, copilot_auth: None, defaulted: false, + query_prefix: String::new(), + document_prefix: String::new(), }) .expect("factory builds compat embedder"); assert_eq!(ok.provider(), "openai-compat"); @@ -101,6 +103,8 @@ async fn factory_builds_compat_embedder_and_requires_base_url() { models_dir: None, copilot_auth: None, defaulted: false, + query_prefix: String::new(), + document_prefix: String::new(), }) { Ok(_) => panic!("compat embedder must not build without a base URL"), Err(err) => err, @@ -119,6 +123,8 @@ async fn factory_builds_compat_embedder_and_requires_base_url() { models_dir: None, copilot_auth: None, defaulted: false, + query_prefix: String::new(), + document_prefix: String::new(), }); assert!( matches!(zero_dim, Err(LlmError::NotConfigured(ref msg)) if msg.contains("greater than zero")), @@ -126,6 +132,156 @@ async fn factory_builds_compat_embedder_and_requires_base_url() { ); } +#[tokio::test] +async fn query_prefix_is_sent_with_embed_query_not_embed_document() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).expect("json body"); + assert_eq!( + body["input"], "query: find the runbook", + "embed_query must send the query prefix ahead of the text" + ); + ResponseTemplate::new(200).set_body_json(embedding_body(8)) + }) + .expect(1) + .mount(&server) + .await; + + let e = OpenAiCompatEmbedder::new(format!("{}/v1", server.uri()), None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + + e.embed_query("find the runbook") + .await + .expect("embed_query succeeds"); +} + +#[tokio::test] +async fn document_prefix_is_sent_with_embed_document_not_embed_query() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).expect("json body"); + assert_eq!( + body["input"], "passage: the runbook says...", + "embed_document must send the document prefix ahead of the text" + ); + ResponseTemplate::new(200).set_body_json(embedding_body(8)) + }) + .expect(1) + .mount(&server) + .await; + + let e = OpenAiCompatEmbedder::new(format!("{}/v1", server.uri()), None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + + e.embed_document("the runbook says...") + .await + .expect("embed_document succeeds"); +} + +#[tokio::test] +async fn unset_prefixes_leave_embed_query_and_embed_document_unchanged() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).expect("json body"); + assert_eq!( + body["input"], "plain text", + "an unset prefix (the default) must not change the request body" + ); + ResponseTemplate::new(200).set_body_json(embedding_body(8)) + }) + .expect(2) + .mount(&server) + .await; + + // No `with_prefixes` call at all — construction alone must be + // byte-identical to before this feature existed. + let e = OpenAiCompatEmbedder::new(format!("{}/v1", server.uri()), None, "nomic-embed-text", 8) + .expect("embedder builds"); + + e.embed_query("plain text") + .await + .expect("embed_query succeeds"); + e.embed_document("plain text") + .await + .expect("embed_document succeeds"); +} + +#[tokio::test] +async fn prefix_is_truncated_together_with_the_text_it_precedes() { + // Truncation must bound prefix+text together; a prefix must never let + // the combined request exceed the server's input cap. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).expect("json body"); + let input = body["input"].as_str().expect("input is a string"); + assert!( + input.starts_with("passage: "), + "prefix must survive truncation: {input:.60}..." + ); + assert!( + input.len() <= 8_000, + "prefix+text must stay within the same hard byte cap as before" + ); + ResponseTemplate::new(200).set_body_json(embedding_body(8)) + }) + .expect(1) + .mount(&server) + .await; + + let e = OpenAiCompatEmbedder::new(format!("{}/v1", server.uri()), None, "nomic-embed-text", 8) + .expect("embedder builds") + .with_prefixes("query: ", "passage: "); + + let long_text = "x".repeat(50_000); + e.embed_document(&long_text) + .await + .expect("embed_document succeeds"); +} + +#[tokio::test] +async fn factory_wires_configured_prefixes_into_the_compat_embedder() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(move |req: &Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).expect("json body"); + assert_eq!(body["input"], "query: via factory"); + ResponseTemplate::new(200).set_body_json(embedding_body(8)) + }) + .expect(1) + .mount(&server) + .await; + + let embedder = build_embedder(EmbedderConfig { + provider: EmbedderChoice::OpenAiCompat, + model: "nomic-embed-text".into(), + dim: 8, + api_key: SecretString::from(String::new()), + base_url: Some(format!("{}/v1", server.uri())), + models_dir: None, + copilot_auth: None, + defaulted: false, + query_prefix: "query: ".into(), + document_prefix: "passage: ".into(), + }) + .expect("factory builds compat embedder with prefixes"); + + embedder + .embed_query("via factory") + .await + .expect("embed_query succeeds"); +} + #[test] fn existing_default_dimension_api_remains_source_compatible() { assert_eq!( diff --git a/crates/ai-memory-mcp/src/admin.rs b/crates/ai-memory-mcp/src/admin.rs index fa00baa19..636a63bea 100644 --- a/crates/ai-memory-mcp/src/admin.rs +++ b/crates/ai-memory-mcp/src/admin.rs @@ -3449,7 +3449,7 @@ async fn handle_lint( .as_ref() .map(|e| ai_memory_consolidate::EmbeddingCoord { provider: e.provider().to_string(), - model: e.model().to_string(), + model: e.model_identity(), dim: e.dim(), }), }, @@ -3599,7 +3599,10 @@ async fn handle_embed( }; let provider = embedder.provider().to_string(); - let model = embedder.model().to_string(); + // Not `.model()`: the purge below must match the identity rows were + // actually stored under (a document-prefix change), not the wire + // model name. See `Embedder::model_identity`. + let model = embedder.model_identity(); let dim = embedder.dim(); let mut totals = EmbedBackfillCounts::default(); @@ -6231,9 +6234,12 @@ async fn copy_purge_merge( let mut src_embeddings: std::collections::HashMap> = std::collections::HashMap::new(); let embed_meta: Option<(String, String, u32)> = if let Some(embedder) = &state.embedder { + // Not `.model()`: only vectors stored under the current document + // identity are "current-model" for the carry-over below to load. + // See `Embedder::model_identity`. let (provider, model, dim) = ( embedder.provider().to_string(), - embedder.model().to_string(), + embedder.model_identity(), embedder.dim(), ); match state diff --git a/crates/ai-memory-mcp/src/server.rs b/crates/ai-memory-mcp/src/server.rs index 3b812f8bb..0f9f7bc81 100644 --- a/crates/ai-memory-mcp/src/server.rs +++ b/crates/ai-memory-mcp/src/server.rs @@ -1964,7 +1964,10 @@ impl AiMemoryServer { let Some(embedder) = &self.embedder else { return None; }; - match embedder.embed(query).await { + // `embed_query`, not `embed`: an asymmetric embedder (configured + // query/document prefixes, or Google's RETRIEVAL_QUERY task type) + // must see this text on the query side, not the document side. + match embedder.embed_query(query).await { Ok(qv) => Some(qv), Err(e) => { tracing::warn!( @@ -1992,7 +1995,11 @@ impl AiMemoryServer { // eligible; with no query vector the vector stream never runs, // so the empty triple is inert rather than a fake identity. let (provider, model, dim) = match (&self.embedder, options.query_vec) { - (Some(e), Some(_)) => (e.provider().to_string(), e.model().to_string(), e.dim()), + // Not `.model()`: eligible stored vectors are keyed by the + // DOCUMENT identity they were embedded under, which a + // configured document prefix folds in. See + // `Embedder::model_identity`. + (Some(e), Some(_)) => (e.provider().to_string(), e.model_identity(), e.dim()), _ => (String::new(), String::new(), 0), }; let fused: Vec<(PageHit, Option)> = if options.explain { @@ -3081,7 +3088,7 @@ impl AiMemoryServer { .as_ref() .map(|e| ai_memory_consolidate::EmbeddingCoord { provider: e.provider().to_string(), - model: e.model().to_string(), + model: e.model_identity(), dim: e.dim(), }), }, @@ -5731,6 +5738,71 @@ mod tests { (tmp, store, server, ws, proj) } + /// A test-only embedder whose three `Embedder` methods each return a + /// different, identifiable vector — mirroring `ai-memory-llm`'s own + /// `health::tests::TaskAwareEmbedder` fixture (same idea, duplicated + /// here because that one is private to its crate). `embed` deliberately + /// returns the SAME vector as `embed_document`: that is what the + /// pre-fix bug actually called (`AiMemoryServer::embed_query` invoked + /// the generic `Embedder::embed` instead of `Embedder::embed_query`), + /// so a regression back to that bug is what this fixture would surface. + struct TaskAwareEmbedder; + + #[async_trait::async_trait] + impl Embedder for TaskAwareEmbedder { + fn provider(&self) -> &'static str { + "task-aware-test" + } + + fn model(&self) -> &str { + "task-aware-test-model" + } + + fn dim(&self) -> u32 { + 2 + } + + async fn embed(&self, _text: &str) -> ai_memory_llm::LlmResult> { + Ok(vec![1.0, 0.0]) + } + + async fn embed_document(&self, _text: &str) -> ai_memory_llm::LlmResult> { + Ok(vec![1.0, 0.0]) + } + + async fn embed_query(&self, _text: &str) -> ai_memory_llm::LlmResult> { + Ok(vec![0.0, 1.0]) + } + } + + /// Regression test for the bug fixed alongside the embedding + /// query/document prefix feature: `AiMemoryServer::embed_query` (the + /// helper `memory_query` calls to vectorize the search query) called + /// `Embedder::embed` instead of `Embedder::embed_query`. For any + /// query/document-asymmetric embedder — Google's task-typed + /// `embedContent`, or the new `openai`/`openai-compat` query/document + /// prefixes — that silently embedded the search query as if it were a + /// document, corrupting vector-stream retrieval. `TaskAwareEmbedder` + /// makes the two paths return distinguishable vectors so the right one + /// being called is a direct, exact-equality assertion rather than an + /// inference from downstream ranking. + #[tokio::test] + async fn memory_query_embeds_the_query_with_embed_query_not_embed() { + let (_tmp, _store, server, _ws, _pj) = setup_server().await; + let server = server.with_embedder(Arc::new(TaskAwareEmbedder)); + + let query_vec = server.embed_query("anything").await; + + assert_eq!( + query_vec.as_deref(), + Some([0.0_f32, 1.0].as_slice()), + "memory_query's embed_query helper must call Embedder::embed_query \ + (query-side vector [0.0, 1.0]), not the generic Embedder::embed \ + (which this fixture deliberately aliases to the document-side \ + vector [1.0, 0.0] to make a regression to the old bug fail loudly)" + ); + } + fn installed_ai_memory_prompt_surface() -> String { let mut prompt = String::from(ai_memory_core::SNIPPET_BODY); for skill in ai_memory_core::routing_skills::MANAGED_SKILLS { diff --git a/crates/ai-memory-wiki/src/wiki.rs b/crates/ai-memory-wiki/src/wiki.rs index 4dac00be7..8c486a03e 100644 --- a/crates/ai-memory-wiki/src/wiki.rs +++ b/crates/ai-memory-wiki/src/wiki.rs @@ -2216,7 +2216,12 @@ impl Wiki { page_id, bytes, embedder.provider().to_string(), - embedder.model().to_string(), + // Not `.model()`: a configured document prefix + // must land under a distinct stored identity so + // a prefix change is never silently mixed with + // vectors embedded under a different (or no) + // prefix. See `Embedder::model_identity`. + embedder.model_identity(), embedder.dim(), ) .await?; @@ -2250,7 +2255,8 @@ impl Wiki { page_id, vector_bytes: f32_vec_to_bytes(&vec), provider: embedder.provider().to_string(), - model: embedder.model().to_string(), + // See the body-embedding write above. + model: embedder.model_identity(), dim: embedder.dim(), }]) .await?; diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 7396f2656..2649845b6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -864,6 +864,12 @@ AI_MEMORY_EMBEDDING_MODEL e.g. text-embedding-3-small, gemini-embedding-001 AI_MEMORY_EMBEDDING_BASE_URL optional override; required for openai-compat AI_MEMORY_EMBEDDING_DIM 1536 (OpenAI, Copilot), 1024 (Voyage), 768 (Google); required explicitly for openai-compat +AI_MEMORY_EMBEDDING_QUERY_PREFIX optional; prepended to query text before + embedding (openai / openai-compat only) +AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX optional; prepended to document text + before embedding (openai / openai-compat + only); e.g. "query: " / "passage: " for + Nemotron-3-Embed / base E5 OPENAI_API_KEY / VOYAGE_API_KEY / GEMINI_API_KEY / GOOGLE_API_KEY LLM_API_KEY accepted for openai with a custom base URL and as optional bearer auth for openai-compat diff --git a/docs/install.md b/docs/install.md index ff774284a..a6dc4c814 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1641,7 +1641,7 @@ If you set only the provider, ai-memory picks a sensible default: | `AI_MEMORY_EMBEDDING_PROVIDER=openai` + `AI_MEMORY_EMBEDDING_BASE_URL=https://api.orcarouter.ai/v1` | `openai/text-embedding-3-small` via [OrcaRouter](https://www.orcarouter.ai) | Uses `EMBEDDING_API_KEY`, else reuses `LLM_API_KEY`, with the OpenAI-compatible embedding client. | | `AI_MEMORY_EMBEDDING_PROVIDER=voyage` | `voyage-3` (1024-dim) | Voyage's current general-purpose recommendation. | | `AI_MEMORY_EMBEDDING_PROVIDER=google` / `gemini` | `gemini-embedding-001` (768-dim) | Google-hosted embeddings via `embedContent`. Set `GEMINI_API_KEY` (or `GOOGLE_API_KEY`). | -| `AI_MEMORY_EMBEDDING_PROVIDER=openai-compat` | no default — set model, dim, and base URL explicitly | Self-hosted engines (Ollama, LM Studio, vLLM). Keyless by default; `EMBEDDING_API_KEY`, else `LLM_API_KEY`, is sent as a bearer token when present (gateways). Example: `AI_MEMORY_EMBEDDING_BASE_URL=http://localhost:11434/v1`, `AI_MEMORY_EMBEDDING_MODEL=nomic-embed-text`, `AI_MEMORY_EMBEDDING_DIM=768`. Switching an existing `openai`+base-URL setup to `openai-compat` changes the stored `{provider, model, dim}` triple — run `ai-memory embed --force` to re-embed. | +| `AI_MEMORY_EMBEDDING_PROVIDER=openai-compat` | no default — set model, dim, and base URL explicitly | Self-hosted engines (Ollama, LM Studio, vLLM). Keyless by default; `EMBEDDING_API_KEY`, else `LLM_API_KEY`, is sent as a bearer token when present (gateways). Example: `AI_MEMORY_EMBEDDING_BASE_URL=http://localhost:11434/v1`, `AI_MEMORY_EMBEDDING_MODEL=nomic-embed-text`, `AI_MEMORY_EMBEDDING_DIM=768`. Switching an existing `openai`+base-URL setup to `openai-compat` changes the stored `{provider, model, dim}` triple — run `ai-memory embed --force` to re-embed. Asymmetric models need `AI_MEMORY_EMBEDDING_QUERY_PREFIX` / `AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX`, e.g. `nvidia/Nemotron-3-Embed-1B-BF16` wants `query: ` / `passage: ` — see [`docs/llm-providers.md`](llm-providers.md) for that and for Qwen3-Embedding/instruction-tuned E5, which need a different (query-only) format. | | `AI_MEMORY_EMBEDDING_PROVIDER=copilot` | `text-embedding-3-small` (1536-dim) | Reuses the `copilot` LLM provider's OAuth login (`ai-memory auth login copilot`, `COPILOT_GITHUB_TOKEN`, or `GITHUB_COPILOT_API_TOKEN`) — no separate API key. Calls Copilot's `/embeddings` endpoint following the OpenAI-compatible contract Copilot documents for chat; that endpoint's exact shape is not covered by a live test against Copilot here, so treat it as needing a real-Copilot smoke test. | > **What we don't recommend:** reasoning-mode models (Claude with extended diff --git a/docs/llm-providers.md b/docs/llm-providers.md index 3d0e69e56..2642018e6 100644 --- a/docs/llm-providers.md +++ b/docs/llm-providers.md @@ -210,6 +210,59 @@ hybrid paths apply the same bounded page-authority adjustment after candidate generation; embeddings improve relevance recall but do not decide which source is canonical. +Asymmetric embedding models — ones trained with a different instruction for +queries than for the text being indexed — need that instruction added to the +text itself; the OpenAI-compatible `/v1/embeddings` wire format has no field +for it. `AI_MEMORY_EMBEDDING_QUERY_PREFIX` and +`AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX` (only applied by the `openai` and +`openai-compat` embedders) are prepended to query and document text +respectively, before the existing truncation, so a long body is still +truncated to the same overall input cap with the prefix included. Both are +empty by default — no behaviour change when unset — and are not trimmed, so +a publisher's trailing space is preserved exactly. For example, NVIDIA's +[`nvidia/Nemotron-3-Embed-1B-BF16`](https://huggingface.co/nvidia/Nemotron-3-Embed-1B-BF16) +(2048-dim, served locally through vLLM/`openai-compat`) specifies +`"query: "` for queries and `"passage: "` for documents per its model card: + +```toml +embedding_provider = "openai-compat" +embedding_model = "nvidia/Nemotron-3-Embed-1B-BF16" +embedding_base_url = "http://localhost:8000/v1" +embedding_dim = 2048 +embedding_query_prefix = "query: " +embedding_document_prefix = "passage: " +``` + +Base E5 models (`intfloat/e5-base-v2`, `e5-large-v2`, multilingual E5, …) +use the same `"query: "` / `"passage: "` convention. Instruction-tuned E5 +variants and Qwen3-Embedding instead need a task-instruction string on the +**query side only** — their documents are embedded plain, with no document +prefix — but the two use **different exact spacing**, confirmed against +each model card: + +- `intfloat/e5-mistral-7b-instruct`: + `embedding_query_prefix = "Instruct: {task description}\nQuery: "` — a + trailing space after `Query:`. +- `Qwen/Qwen3-Embedding-0.6B` (and the other Qwen3-Embedding sizes): + `embedding_query_prefix = "Instruct: {task description}\nQuery:"` — **no** + trailing space; the query text follows the colon directly. + +Fill in your own task description for `{task description}`, leave +`embedding_document_prefix` unset for both, and don't copy one model's +exact string for the other — the trailing-space difference is +publisher-specified, not a typo. + +Changing `embedding_query_prefix` alone never requires re-embedding: only +`embedding_document_prefix` is folded into the stored embedding identity +(a non-empty value changes what `ai-memory` treats as the current +`{provider, model, dim}` triple for that page's vectors), so a document +prefix change makes `memory_query` stop matching the old vectors and the +next `ai-memory embed` (scheduled or manual) pass re-embed them +automatically — no `--force` needed. `--force` remains useful to +re-embed everything immediately rather than waiting for scheduled +backfill, or to force a re-embed for a reason the triple alone doesn't +capture (e.g. touching up the vLLM server's tokenizer or quantization). + `AI_MEMORY_EMBEDDING_PROVIDER=copilot` reuses the same Copilot OAuth login as the `copilot` LLM provider (`ai-memory auth login copilot` or `COPILOT_GITHUB_TOKEN`/`GITHUB_COPILOT_API_TOKEN`) — no separate API key.