Skip to content

feat(embedding): query/document prefixes for OpenAI-compatible embedders - #859

Merged
akitaonrails merged 3 commits into
akitaonrails:release/2.5from
seathatflowsinourveins:embedding-query-document-prefix
Sep 23, 2026
Merged

akitaonrails merged 3 commits into
akitaonrails:release/2.5from
seathatflowsinourveins:embedding-query-document-prefix

Conversation

@seathatflowsinourveins

Copy link
Copy Markdown
Contributor

What changed

Adds two optional config keys, embedding_query_prefix and
embedding_document_prefix (env: AI_MEMORY_EMBEDDING_QUERY_PREFIX /
AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX), applied by the openai and
openai-compat embedders. Each is prepended to the text before the
existing 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 a config.toml value rather than being treated as unset.
OpenAiEmbedder/OpenAiCompatEmbedder gained a with_prefixes builder
and now override embed_query/embed_document; embed is unchanged.

Also fixes memory_query's internal embed_query helper
(ai-memory-mcp/src/server.rs): it called the generic Embedder::embed
instead of Embedder::embed_query, so a query/document-asymmetric
embedder (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_prefix into the
stored embedding identity (Embedder::model_identity(), default
model().to_string() — every other provider is unaffected), via a
versioned SHA-256 fingerprint ("+dp1-" + first 16 hex chars, not a
32-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 own
commit (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.ps1 notes that an
operator's own $env:NAME = '' additionally needs PowerShell 7.5+ for
that to leave a set-but-empty variable rather than deleting it; not
exercised on a real pwsh runtime). 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/embeddings wire format has no field for it, so the client has to
add it. nvidia/Nemotron-3-Embed-1B-BF16 and base E5 models specify a
simple "query: " / "passage: " pair. Instruction-tuned E5 variants and
Qwen3-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 of
these through openai-compat got degraded retrieval with no config-level
fix; only google had built-in query/document asymmetry.

Test plan

  • cargo fmt --all -- --check and git diff --check pass
  • cargo clippy --workspace --all-targets -- -D warnings passes, 0
    warnings
  • 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 to
    embeddings.
  • On a v2.4.0 + this patch build: ai-memory-llm 270 passed/0
    failed, ai-memory-mcp 465 passed/0 failed — identical to the
    feature branch.
  • Unit, wiremock and loader tests cover: prefix applied per side, unset is
    a true no-op, model_identity behaves correctly per provider, the
    fingerprint does not collide on a pair that collided under an earlier
    32-bit-hash design, Config::load through a TOML file in an isolated
    child process with both prefix env vars removed, memory_query calling
    embed_query not embed, and fake-Docker wrapper argument capture
    across unset/empty/whitespace-only/non-empty.
  • ai-memory-consolidate's embed_backfill.rs suite is extended with two
    tests: 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

  • Verify the name/email on every commit before requesting merge (see
    CONTRIBUTING.md).

Release impact

  • Minor — additive: two new config keys / env vars, no default
    changed. The 3rd commit changes what gets stored only for a
    non-default (non-empty document-prefix) configuration.

CHANGELOG (merge gate)

  • CHANGELOG.md [Unreleased] → ### Added entry for the new keys
    (including the identity-folding behaviour) and a ### Fixed entry
    for the embed→embed_query bug.

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

…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>
@akitaonrails
akitaonrails changed the base branch from main to release/2.5 September 23, 2026 17:09
@akitaonrails
akitaonrails merged commit 7552d4f into akitaonrails:release/2.5 Sep 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants