Repository navigation
feat(embedding): query/document prefixes for OpenAI-compatible embedders - #859
Merged
akitaonrails merged 3 commits intoSep 23, 2026
Conversation
…nai-compat embedders Adds `embedding_query_prefix` / `embedding_document_prefix` config keys (env: AI_MEMORY_EMBEDDING_QUERY_PREFIX / AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX) applied by 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 (Nemotron-3-Embed, base E5, Qwen3-Embedding) need a "query: " / "passage: " instruction their publisher specifies; the OpenAI-compatible /v1/embeddings wire format has no field for it. Empty by default, and not trimmed so a publisher's trailing space is preserved. Also fixes memory_query's embed_query helper (ai-memory-mcp/src/server.rs), found while auditing every embed/embed_query/embed_document call site: it called the generic Embedder::embed instead of Embedder::embed_query, so an asymmetric embedder (google's task-typed embeddings, or these new prefixes) embedded the search query on the document side. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ssion tests
figment's Env provider parses each var's string with its loose-value
parser, whose bare/unquoted branch calls .trim() (figment 0.10.19's
src/value/parse.rs:78, verified against the vendored source) — so
AI_MEMORY_EMBEDDING_QUERY_PREFIX="query: " reached Config::embedding_query_prefix
as "query:", silently dropping the publisher-significant trailing space.
Config::load now overlays the raw (untrimmed) env value for the two prefix
keys via figment::providers::Serialized (which hands figment an
already-typed value, bypassing the string parser), merged after the
generic Env::prefixed pass so it still wins over a config.toml value.
Presence, not non-emptiness, is the signal: a variable set to "" is a
deliberate override clearing a config.toml-configured prefix, distinct
from the variable being absent. Both wrapper scripts (bin/ai-memory,
bin/ai-memory.ps1) now forward these two keys on presence for the same
reason; a non-empty override was already forwarded correctly, only the
empty-override case was silently dropped by the [ -n ] check every other
forwarded var correctly uses.
The overlay itself lives in overlay_embedding_prefixes(), a pure function
taking the env values as parameters rather than reading std::env::var
itself, so it stays directly unit-testable without mutating process
environment or the current directory — the same pattern
ai-memory-cli/src/commands/path_util.rs's agent_config_home and
ai-memory-hooks's drain_with_live_token already use. std::env::set_var is
unsafe under edition 2024 and forbidden workspace-wide, and
figment::Jail calls std::env::set_current_dir on the real process
internally, racing every other test in this crate's multi-threaded lib
test binary that relies on cwd. Four pure unit tests build a minimal
in-memory Figment and assert on the merged Config directly. One
additional test exercises the real Config::load end to end through a
TOML file and an explicit absolute path; since this process's own
environment is shared with every other test in the binary and could
already carry one of the two prefix vars from the test runner's shell,
that test re-execs this same test binary filtered to just itself as a
genuinely separate child process, with both vars removed via
Command::env_remove — real process isolation rather than an in-process
assumption about the ambient environment.
Plus: a wiremock transport test proving OpenAiEmbedder (not just
OpenAiCompatEmbedder) sends the configured prefix on the wire; a
regression test for the memory_query embed_query fix from the prior
commit, using a task-aware fixture embedder whose
embed/embed_document/embed_query methods return distinguishable vectors
so a regression back to calling the wrong one fails a direct equality
assertion rather than an inferred ranking change; fake-Docker argument
tests for the wrapper's presence-based forwarding across
unset/empty/whitespace-only/non-empty, with both env vars explicitly
removed from the child environment before each case so the "unset" case
cannot silently inherit an ambient export from the test runner.
bin/ai-memory.ps1 also notes, briefly, the PowerShell/.NET version an
operator needs for `$env:NAME = ''` to reach the wrapper as set-but-empty
rather than deleted — no pwsh runtime was available to exercise it live.
Also narrows the query/document-prefix docs (config.rs, docs/llm-providers.md,
CHANGELOG.md) with exact, separate templates: Nemotron-3-Embed
(nvidia/Nemotron-3-Embed-1B-BF16) and base E5 use "query: "/"passage: ";
e5-mistral-7b-instruct wants "Instruct: {task}\nQuery: " (trailing space);
Qwen3-Embedding wants "Instruct: {task}\nQuery:" (no trailing space) — both
leave documents plain.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
A document prefix change previously left old vectors silently searchable under the same (provider, model, dim) triple, because the stored/matched identity was always the raw wire model name regardless of the configured prefix. Adds Embedder::model_identity() (default: model().to_string(), so every provider without a document-prefix concept is unaffected byte-for-byte). OpenAiEmbedder/OpenAiCompatEmbedder override it: an empty document_prefix reproduces model() exactly (the legacy identity — upgrading installs with no prefix configured need no migration); a non-empty one appends a versioned SHA-256 fingerprint (document_prefix_identity(): "+dp1-" plus the first 16 hex characters of SHA-256(document_prefix)) so two different prefixes — or a prefix vs. none — never collide. A 32-bit fnv1a fingerprint (reusing the hash already in this module for SyntheticEmbedder) collides on same-length ASCII input — "Document category 5hvhw0: " and "Document category 1i4yh8i: " both fingerprint to 5c7f37a2, which would silently mix two different prefixes' vectors under one identity — so this uses SHA-256 instead; that pair is now a regression test. Only the document prefix feeds this; the query prefix never does, since a query is never stored. Threaded through every place that builds or matches the stored (provider, model, dim) triple: body and abstract page writes (wiki.rs), backfill's write and stale-row detection (ai-memory-consolidate's run_embedding_backfill, shared by the scheduled tick and POST /admin/embed per its own module doc), retrieval's eligibility filter (AiMemoryServer::search_project), the A3 cold-cluster dedup coordinate and the zero-LLM contradiction lint's EmbeddingCoord (admin.rs, server.rs, serve.rs), the startup refuse-on-mismatch check and its diagnostic log (serve.rs), admin.rs's reembed purge and its bulk-page-move carry-over, and the HealthRecordingEmbedder wrapper every configured embedder is actually wrapped in before reaching any of the above (missed, this would have silently defeated the whole feature). Net effect: changing embedding_document_prefix makes memory_query stop matching existing pages' vectors and the next backfill pass (scheduled or manual) re-embed them automatically — no `ai-memory embed --force` needed for that case specifically, though it remains available. Changing only embedding_query_prefix still needs no rebuild at all. CHANGELOG.md and docs/llm-providers.md are updated with this re-embed-automatically behaviour. Tests: unit coverage for document_prefix_identity (including the fnv32 collision pair as an explicit regression case) and model_identity on both embedders and the health wrapper; two ai-memory-consolidate integration tests exercising the real flow — a document-prefix change making a normal backfill pass treat a page as stale and re-embed it (replacing, not duplicating, the page_embeddings row, since it's PRIMARY KEY page_id), and a POST /admin/embed-style forced rebuild (purge then re-embed) with one page failing mid-run, asserting the purge still covers both pages and the failed page is left with no row rather than a stale or half-written one. Kept as its own commit, on top of a complete commit 2, so it can be reviewed and dropped independently of the rest of this PR if the documented-manual-`--force` alternative is preferred; the rest of the PR does not depend on it. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
seathatflowsinourveins
added a commit
to native-agent-stack/native-agent-stack
that referenced
this pull request
Sep 23, 2026
…rement, live deployment) (#127) * Record the local memory-embedder evidence: upstream-aligned benchmark, hybrid re-measurement, live deployment Adds evidence/artifacts/local-model-alignment-20260923/ (hash-listed, no selection labels): - memory-embedder-dense-benchmark.json: current embedders set up per their model cards on a sealed 130-page memory corpus with a two-annotator query set; Nemotron-3-Embed-1B highest (nDCG@10 0.858 vs the all-MiniLM-L6-v2 reference 0.713; its lead over Qwen3-Embedding and granite-r2 is within noise); - memory-embedder-hybrid-retrieval.json: through ai-memory's own memory_query, Nemotron with its query/passage prefixes at 8192 context beat the MiniLM configuration by nDCG@10 +0.078 (95% CI 0.037-0.128) against a preregistered 0.03 bar; 4096 context loses nine pages; - ai-memory-embedder-deployment.json: the host's live switch, the patched 2.3.2 build (upstream PR akitaonrails/ai-memory#859, binary sha256), configuration, observed results and rollback. Each carries the independent verifier's findings, including the weaker lock and log provenance. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * Apply the evidence review corrections: q29/q30 note, Q1 annotator limit, observed-vs-logged labels, measurement basis, rollback not exercised Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: seathatflowsinourveins <234074349+seathatflowsinourveins@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This was referenced Sep 23, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
Adds two optional config keys,
embedding_query_prefixandembedding_document_prefix(env:AI_MEMORY_EMBEDDING_QUERY_PREFIX/AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX), applied by theopenaiandopenai-compatembedders. Each is prepended to the text before theexisting truncation. Both default to empty (no behaviour change when
unset) and are never trimmed: a publisher's trailing space (e.g.
"query: ") is part of the instruction, and an env var explicitly set to""clears aconfig.tomlvalue rather than being treated as unset.OpenAiEmbedder/OpenAiCompatEmbeddergained awith_prefixesbuilderand now override
embed_query/embed_document;embedis unchanged.Also fixes
memory_query's internalembed_queryhelper(
ai-memory-mcp/src/server.rs): it called the genericEmbedder::embedinstead of
Embedder::embed_query, so a query/document-asymmetricembedder (Google's task-typed embeddings, or a configured query prefix)
embedded search queries on the document side.
A third change folds a non-empty
embedding_document_prefixinto thestored embedding identity (
Embedder::model_identity(), defaultmodel().to_string()— every other provider is unaffected), via aversioned SHA-256 fingerprint (
"+dp1-" + first 16 hex chars, not a32-bit hash, to avoid collisions between different prefixes). A document
prefix change makes retrieval, backfill, and cleanup treat existing pages
as stale and re-embed them automatically; an empty prefix keeps the exact
pre-existing identity, so upgrading installs need no migration. Real
surface area across 8 production Rust files for a feature whose
alternative is a documented
ai-memory embed --force, so it is its owncommit (3rd) on a complete 2nd commit — verified to drop cleanly.
Both wrapper scripts forward the two new env vars on presence rather than
non-emptiness, so a present-but-empty override reaches the container and
clears a TOML-configured prefix (
bin/ai-memory.ps1notes that anoperator's own
$env:NAME = ''additionally needs PowerShell 7.5+ forthat to leave a set-but-empty variable rather than deleting it; not
exercised on a real
pwshruntime). Docs updated:docs/llm-providers.md,docs/install.md,docs/ARCHITECTURE.md.Why
Asymmetric embedding models use a different instruction for the query
side than for the indexed-document side. The OpenAI-compatible
/v1/embeddingswire format has no field for it, so the client has toadd it.
nvidia/Nemotron-3-Embed-1B-BF16and base E5 models specify asimple
"query: "/"passage: "pair. Instruction-tuned E5 variants andQwen3-Embedding instead need a task-instruction string on the query side
only, with a publisher-specific trailing-space convention that differs
between the two (
intfloat/e5-mistral-7b-instruct:"Instruct: {task}\nQuery: "; Qwen3-Embedding:"Instruct: {task}\nQuery:"). Without this, an operator running any ofthese through
openai-compatgot degraded retrieval with no config-levelfix; only
googlehad built-in query/document asymmetry.Test plan
cargo fmt --all -- --checkandgit diff --checkpasscargo clippy --workspace --all-targets -- -D warningspasses, 0warnings
cargo test --workspace --all-targets --no-fail-fast: 3546 passed,3 failed, 14 ignored. The 3 failures also fail, with identical
names and root causes, on unmodified
main, unrelated toembeddings.
ai-memory-llm270 passed/0failed,
ai-memory-mcp465 passed/0 failed — identical to thefeature branch.
a true no-op,
model_identitybehaves correctly per provider, thefingerprint does not collide on a pair that collided under an earlier
32-bit-hash design,
Config::loadthrough a TOML file in an isolatedchild process with both prefix env vars removed,
memory_querycallingembed_querynotembed, and fake-Docker wrapper argument captureacross unset/empty/whitespace-only/non-empty.
ai-memory-consolidate'sembed_backfill.rssuite is extended with twotests: a document-prefix change makes backfill treat a page as stale
and re-embed it; a forced rebuild with one page failing mid-run still
purges both pages and leaves the failed page with no row.
Commit attribution
CONTRIBUTING.md).
Release impact
changed. The 3rd commit changes what gets stored only for a
non-default (non-empty document-prefix) configuration.
CHANGELOG (merge gate)
CHANGELOG.md[Unreleased]→### Addedentry for the new keys(including the identity-folding behaviour) and a
### Fixedentryfor the
embed→embed_querybug.Notes for reviewers
The 3rd commit is the one design decision worth an explicit call — see
"What changed" above; it's kept separate specifically so it can be
reviewed and dropped on its own.
🤖 Generated with Claude Code