From f134c50d390d22a29ae7db02a42ecfb4e67c57d7 Mon Sep 17 00:00:00 2001 From: Julian lechner Date: Thu, 1 Oct 2026 13:13:06 +0200 Subject: [PATCH] fix(retrieval): match the Qwen query prompt format Drop the trailing space after "Query:" as in the model card and include the prompt in the query cache key. Co-Authored-By: Claude Opus 5.5 --- scripts/tests/test-retrieval-server.py | 4 ++++ shared/retrieval/server.py | 7 ++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/scripts/tests/test-retrieval-server.py b/scripts/tests/test-retrieval-server.py index 7052bea..45f2f7c 100644 --- a/scripts/tests/test-retrieval-server.py +++ b/scripts/tests/test-retrieval-server.py @@ -133,6 +133,10 @@ class RetrievalServerTests(unittest.TestCase): self.assertNotEqual(MODULE.daemon_address(data_dir), old_address) self.assertNotEqual(MODULE.daemon_lock_path(data_dir), old_lock) + def test_query_prompt_matches_qwen_format(self): + self.assertTrue(MODULE.QUERY_PROMPT.startswith('Instruct: ')) + self.assertTrue(MODULE.QUERY_PROMPT.endswith('\nQuery:')) + def test_preferred_device_uses_cuda_with_cpu_fallback(self): with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))): self.assertEqual(MODULE.preferred_device(), 'cuda') diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 4a99c67..80f068a 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -36,6 +36,8 @@ CLI_TIMEOUT_SECONDS = 300 RRF_K = 60 INDEX_VERSION = 3 RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question." +# Qwen3-Embedding query format from the model card: no space after "Query:". +QUERY_PROMPT = f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery:" TEXT_EXTENSIONS = {".c", ".cjs", ".cpp", ".cs", ".csproj", ".css", ".go", ".gradle", ".html", ".java", ".js", ".json", ".jsx", ".kt", ".kts", ".md", ".mjs", ".props", ".ps1", ".psm1", ".py", ".razor", ".rs", ".scss", ".sh", ".sql", ".targets", ".toml", ".ts", ".tsx", ".txt", ".xaml", ".xml", ".yaml", ".yml"} @@ -201,8 +203,7 @@ def shared_encoder(): if _shared_encoder is None: from sentence_transformers import SentenceTransformer _shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS, - prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "}, - device=preferred_device()) + prompts={"query": QUERY_PROMPT}, device=preferred_device()) return _shared_encoder @@ -522,7 +523,7 @@ class Index: def cache_key(self, query: str, top_k: int): identity = json.dumps([ query, top_k, - RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, + RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, QUERY_PROMPT, DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K, ], ensure_ascii=False, separators=(",", ":")) return hashlib.sha256(identity.encode("utf-8")).hexdigest()