Skip to content

fix:hybrid retriever: reuse query embedding across retries - #38

Open
KeyCOrigin wants to merge 2 commits into
OpenDCAI:mainfrom
KeyCOrigin:dev
Open

KeyCOrigin wants to merge 2 commits into
OpenDCAI:mainfrom
KeyCOrigin:dev

Conversation

@KeyCOrigin

Copy link
Copy Markdown

Summary

In HybridRetriever._retrieve, the vector-search step is wrapped by
recover_read and can be retried when the vector store read fails. Before
this change, every retry called embed_query(query) again, even though the
query text never changes across retries.

This PR caches the query embedding for the duration of a single retrieval:

  • Compute embed_query(query) once, before the first vector query.
  • Reuse the same embedding on subsequent vector-search retries.
  • If the embedding step itself fails, the cache stays empty and the next
    attempt re-embeds (no stale value is reused).

Changes

  • datamind/capabilities/kb/providers/hybrid_retriever.py
    • Added a per-retrieval query_vector local that holds the completed
      embedding across recover_read retries.
    • Moved stage = "kb.embedding" to run only when the embedding is
      actually computed.
    • Kept the store query unchanged otherwise (same filters, top-k,
      versioned reads, and fallback behaviour).

Notes

  • Scope is intentionally limited to a single retrieval call; no
    cross-request embedding cache is added here.
  • No behaviour change for the happy path or for the BM25 fallback path.

Verification

  • python3 -m py_compile datamind/capabilities/kb/providers/hybrid_retriever.py
    passes.

This branch has not been deployed

No deployments
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.

1 participant