The retrieval worked and returned the wrong file

TL;DR#

Embedding search over a codebase returns what is semantically similar to the query. For code, the most similar thing to PaymentProcessor.kt is PaymentProcessorTest.kt — same vocabulary, same type names, same domain.

So a query about payment processing retrieves the test, the mock, the fixture, and the deprecated copy before it retrieves the implementation. The retrieval is working exactly as specified; the specification was wrong, because similarity is not the same relation as relevance.

The fix is not a better embedding model. It is to combine a structural signal — the call graph, the file's role, whether it is referenced by the thing you are changing — with the semantic one.

The case this post is about#

A code-assistant retrieval step, built the standard way:

def retrieve(query: str, k: int = 8) -> list[Chunk]:
    q = embed(query)
    return vector_db.search(q, limit=k)       # cosine similarity over 400-line chunks

Index: every file in the repository, chunked at 400 lines, embedded with a general-purpose model. About 14,000 chunks.

The symptom#

Query: "how does the retry policy decide the backoff delay"

rank  score   file
1     0.891   core/retry/RetryPolicyTest.kt          ← the test
2     0.887   core/retry/FakeRetryPolicy.kt          ← the fake
3     0.884   docs/archive/RETRY_DESIGN_V1.md        ← superseded design doc
4     0.879   core/retry/RetryPolicyTest.kt (chunk 2)
5     0.876   legacy/retry/RetryPolicy.kt            ← deleted-but-not-deleted copy
6     0.871   core/retry/RetryPolicy.kt              ← the answer, rank 6
7     0.869   core/sync/SyncRetryWrapper.kt
8     0.864   core/retry/RetryPolicyTest.kt (chunk 3)

The implementation at rank 6, behind three chunks of its own test, a fake, an archived design document and a legacy copy.

Three consequences:

Answers derived from the test. The assistant explained the backoff as "jittered between 0 and the ceiling, per the test's assertion of the distribution" — true of the test's expectation and not a description of the code.

Answers derived from the archived doc. RETRY_DESIGN_V1.md described exponential backoff without jitter, which had been the design two years earlier. The assistant confidently described behaviour that had been changed deliberately.

And edits applied to the legacy copy. legacy/retry/RetryPolicy.kt scored above the real one, so a "fix the backoff ceiling" request modified a file that nothing imports.

Measured over 50 real queries:

queries where the top-3 contained the primary implementation:   18/50  (36%)
queries where the top-3 contained a test of it:                 44/50  (88%)
queries where the top-1 was a test, fake or fixture:            31/50  (62%)

Why the obvious fix failed#

"Use a code-specific embedding model." We switched to one trained on code. The implementation's average rank improved from 5.8 to 4.1 — real, and nowhere near enough, because the reason a test is similar to its subject is not a weakness of the model. They genuinely are similar.

"Exclude test files from the index." Our first real fix, and it broke the queries where the test is the answer: "how is RetryPolicy supposed to behave on a 429" is answered best by the test. Exclusion trades one failure for another.

"Increase k to 20." The implementation now appears, buried in 20 chunks of mostly noise — and 20 chunks is 8k tokens, which puts us back in the context-budget problem with a worse signal-to-noise ratio.

"Add -test to the query." Query manipulation, brittle, and it does not address the archived doc or the legacy copy, which are not tests.

The actual mechanism#

Embedding search ranks by semantic similarity. What you want is relevance to the task, and for code those two relations differ in specific, predictable ways:

A test shares its subject's entire vocabulary. Type names, method names, domain terms, often whole expressions. By construction it is near-maximally similar, and it describes expectations rather than behaviour.

A fake or mock is similar and deliberately wrong. It implements the same interface with simplified behaviour — the one thing you must not read as the real logic.

A superseded document is similar and false. It was written about this exact subject, which is why it scores well, and it describes a state of the world that no longer exists.

The common structure is that every one of these is similar because it is wrong. A test, a fake and a stale design doc are all derived from the real thing, so they inherit its vocabulary while inverting its authority. That is not a ranking bug to be tuned away — it is the similarity metric working correctly on a question nobody asked it.

Which gives the sharp version: retrieval quality cannot be improved by making the embeddings better, because the failure is not in the distance calculation. OrderServiceTest.kt genuinely is one of the nearest neighbours of "how do we calculate order totals". Any model that ranked it lower would be worse at its actual job.

So the lever is what is in the index and what metadata travels with each chunk, not the model. A corpus that mixes specification, implementation, expectation and simulation with no way to distinguish them is asking the retriever to make a judgement it has no input for — and the fix below is entirely about supplying that input rather than about search quality at all.

A duplicate or legacy copy is nearly identical. Similarity cannot distinguish the live one from the dead one, because the distinguishing fact is structural: whether anything imports it.

And that is the key insight: every one of those distinctions is structural, not semantic.

is it referenced by the code I am changing?      call graph
is it the live implementation or a dead copy?    import graph
is it a test, a fake, a doc, or production?      file role
has it changed recently?                         git history

Similarity answers "what talks about this". Relevance answers "what decides this". For code the second question is largely answerable from the graph, which is deterministic and free.

The fix#

Hybrid retrieval: structural candidates first, semantic ranking within them, and role-aware weighting.

@dataclass
class Candidate:
    chunk: Chunk
    semantic: float          # cosine similarity
    structural: float        # graph distance from the task's focus
    role: Role               # PRODUCTION | TEST | FAKE | DOC | DEAD

def retrieve(query: str, focus: Path | None, k: int = 8) -> list[Chunk]:
    # 1. STRUCTURAL candidates — deterministic, from the graph
    structural = set()
    if focus:
        structural |= imports_of(focus) | importers_of(focus)      # 1 hop
        structural |= siblings_in_module(focus)

    # 2. SEMANTIC candidates
    semantic = {c.path for c in vector_db.search(embed(query), limit=40)}

    # 3. score, with role weights and a hard filter on dead code
    scored = []
    for path in structural | semantic:
        role = classify(path)
        if role is Role.DEAD:                 # nothing imports it: never return it
            continue
        scored.append(Candidate(
            chunk=chunk_for(path),
            semantic=cosine(embed(query), embed_of(path)),
            structural=1.0 if path in structural else 0.0,
            role=role,
        ))
    return [c.chunk for c in sorted(scored, key=rank, reverse=True)[:k]]

ROLE_WEIGHT = {Role.PRODUCTION: 1.00, Role.TEST: 0.55, Role.DOC: 0.50, Role.FAKE: 0.25}

def rank(c: Candidate) -> float:
    return (0.55 * c.semantic + 0.45 * c.structural) * ROLE_WEIGHT[c.role]

Four decisions:

Dead code is filtered, not down-weighted. legacy/retry/RetryPolicy.kt has zero importers, so it is unreachable and must never be returned — a deterministic graph fact, not a judgement:

def classify(path: Path) -> Role:
    if not importers_of(path) and not is_entry_point(path): return Role.DEAD
    if path.match("**/test/**") or path.name.endswith("Test.kt"): return Role.TEST
    if path.name.startswith(("Fake", "Mock", "Stub")): return Role.FAKE
    if path.suffix == ".md": return Role.DOC
    return Role.PRODUCTION

Roles are down-weighted, not excluded. A test at 0.55 still surfaces when it is genuinely the best match — which is the case the exclusion attempt broke. The fake at 0.25 effectively never wins unless nothing else matches.

Structural candidates are injected regardless of their embedding score. One hop of the import graph from the file being changed is almost always relevant, and sometimes scores poorly because it uses different vocabulary.

And archived docs get the DEAD treatment by path convention:

DEAD_PATH_MARKERS = ("docs/archive/", "/deprecated/", "/legacy/", ".old.")

Which is the cheap version of a real staleness signal — and it is what stopped the two-year-old design doc from being cited as current behaviour.

How we proved it#

The 50-query benchmark, with the primary implementation hand-labelled per query:

                                  baseline   code-embeddings   hybrid
implementation in top-3              36%           48%           88%
implementation at rank 1             14%           22%           66%
top-1 was a test/fake/fixture        62%           51%            8%
dead-code chunk anywhere in top-8    34%           31%            0%
mean rank of implementation          5.8           4.1           1.6

The 0% on dead code is the filter, not the ranking — that is the row I trust most, because it is deterministic rather than statistical.

That the test-is-the-answer case still works — the thing exclusion broke:

7 queries where a test is the correct primary source
  baseline:   test in top-3   7/7
  exclusion:  test in top-3   0/7      ← the fix that was worse
  hybrid:     test in top-3   7/7

And the downstream effect, which is what any of this is for:

50 queries, answer graded against the actual code by hand
  baseline answers correct:    28/50
  hybrid answers correct:      44/50

edits applied to the wrong file
  baseline:  6 of 20 edit requests
  hybrid:    0 of 20

Plus the cost, since structural lookups are not free:

per retrieval
  baseline:   1 embedding call,  38ms
  hybrid:     1 embedding call, 112ms   (graph queries from a pre-built index)

The import graph is built once per commit and cached, so the 74ms is index lookups rather than parsing.

Two honest limits. classify is convention-based — a fake not named Fake*, or a live file with no importers because it is reached by reflection, will be misclassified. We found two reflection-reached files and added them to an explicit entry-point list, which is a maintenance cost rather than a solved problem. And the 0.55/0.45 weighting was tuned on these same 50 queries, so it is fitted to them; the role weights transferred to a second repository, the blend ratio needed retuning.

What this generalizes to#

Similarity is a proxy for relevance, and proxies fail in structured domains. Code has an import graph, documents have citations, products have a taxonomy. Where real structure exists, using only the embedding discards your best signal.

Some distinctions are not semantic and never will be. Live versus dead, real versus fake, current versus superseded. No embedding model can tell these apart, because the difference is not in the text. Stop trying to solve them with a better model.

Down-weight rather than exclude. Removing a whole class of document breaks the queries where that class is the answer. A weight preserves the exception; a filter destroys it — reserve filters for facts, like unreachability.

Hand-label a benchmark before tuning. Fifty queries with a known correct answer took an afternoon and turned every subsequent decision into a measurement. Without it, "the retrieval seems better" is the only available verdict, and it is how a 36% top-3 rate survives.

Related#

• Context is a budget, and what you load first wins — why retrieval was attractive
• Your agent's tool output is untrusted input — what retrieved content can carry
• The most valuable pull request that quarter deleted 4,000 lines — why the legacy copy was there at all

Keep reading#

This is part of Context engineering for code — what to put in a model's context, measured rather than assumed. Follow MobileByteSensei for the rest.