feat(qwen): rerank 50 candidates and return five
This commit is contained in:
1 parent
e97404a0a1
commit
0f8aa0a25d
2 files changed
+49
-4
No files matched your search
@@ -205,6 +205,51 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4])
|
self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4])
|
||||||
self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2])
|
self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2])
|
||||||
|
|
||||||
|
def test_hybrid_candidates_and_document_metadata(self):
|
||||||
|
with TemporaryDirectory() as temporary:
|
||||||
|
root = Path(temporary)
|
||||||
|
for number in range(110):
|
||||||
|
(root / f'{number:03}.md').write_text(f'# Symbol{number}\ncontent {number}', encoding='utf-8', newline='\n')
|
||||||
|
index = MODULE.Index(root, root / 'data')
|
||||||
|
documents = []
|
||||||
|
|
||||||
|
def encode(values, **kwargs):
|
||||||
|
if kwargs.get('prompt_name') != 'query':
|
||||||
|
documents.extend(values)
|
||||||
|
return [embedding() for _ in values]
|
||||||
|
|
||||||
|
pairs = []
|
||||||
|
index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
|
||||||
|
index.reranker = lambda: SimpleNamespace(predict=lambda values, **kwargs:
|
||||||
|
pairs.extend(values) or [float('109.md' in document) for _, document in values])
|
||||||
|
numpy = SimpleNamespace(dot=lambda a, b: 1.0,
|
||||||
|
frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
|
||||||
|
original_fusion = MODULE.reciprocal_rank_fusion
|
||||||
|
rankings = []
|
||||||
|
|
||||||
|
def fuse(*values):
|
||||||
|
rankings.extend(values)
|
||||||
|
return original_fusion(*values)
|
||||||
|
|
||||||
|
try:
|
||||||
|
with patch.dict('sys.modules', numpy=numpy), patch.object(MODULE, 'reciprocal_rank_fusion', fuse):
|
||||||
|
results = index.search('109.md Symbol109 content', 20)
|
||||||
|
self.assertEqual([len(ranking) for ranking in rankings], [50, 50])
|
||||||
|
self.assertEqual(len(pairs), 50)
|
||||||
|
self.assertEqual(len(results), 5)
|
||||||
|
self.assertEqual(results[0]['path'], '109.md')
|
||||||
|
self.assertEqual(results[0]['symbol'], 'Symbol109')
|
||||||
|
self.assertIn('Path: 109.md\nSymbol: Symbol109\n\n# Symbol109\ncontent 109', documents)
|
||||||
|
self.assertTrue(all(document in documents for _, document in pairs))
|
||||||
|
with patch.dict('sys.modules', numpy=numpy):
|
||||||
|
self.assertEqual(len(index.search('" OR * () : -')), 5)
|
||||||
|
self.assertEqual(len(index.search('***')), 5)
|
||||||
|
(root / '109.md').unlink()
|
||||||
|
index.search('Symbol109')
|
||||||
|
self.assertEqual(index.db.execute('select count(*) from lexical where lexical match ?', ('"Symbol109"',)).fetchone()[0], 0)
|
||||||
|
finally:
|
||||||
|
index.db.close()
|
||||||
|
|
||||||
def test_old_cache_is_preserved_but_not_reused(self):
|
def test_old_cache_is_preserved_but_not_reused(self):
|
||||||
with TemporaryDirectory() as temporary:
|
with TemporaryDirectory() as temporary:
|
||||||
root = Path(temporary)
|
root = Path(temporary)
|
||||||
|
|||||||
@@ -22,8 +22,8 @@ CHUNK_TOKENS = 400
|
|||||||
OVERLAP_TOKENS = 64
|
OVERLAP_TOKENS = 64
|
||||||
DENSE_CANDIDATES = 50
|
DENSE_CANDIDATES = 50
|
||||||
LEXICAL_CANDIDATES = 50
|
LEXICAL_CANDIDATES = 50
|
||||||
RERANK_CANDIDATES = 20
|
RERANK_CANDIDATES = 50
|
||||||
FINAL_RESULTS = 20
|
FINAL_RESULTS = 5
|
||||||
RRF_K = 60
|
RRF_K = 60
|
||||||
INDEX_VERSION = 2
|
INDEX_VERSION = 2
|
||||||
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
||||||
@@ -216,7 +216,7 @@ class Index:
|
|||||||
self.delete_path(relative)
|
self.delete_path(relative)
|
||||||
self.db.execute("delete from files where path = ?", (relative,))
|
self.db.execute("delete from files where path = ?", (relative,))
|
||||||
|
|
||||||
def search(self, query: str, top_k: int = 5):
|
def search(self, query: str, top_k: int = FINAL_RESULTS):
|
||||||
if not query.strip():
|
if not query.strip():
|
||||||
raise ValueError("query must not be empty")
|
raise ValueError("query must not be empty")
|
||||||
with self.lock:
|
with self.lock:
|
||||||
@@ -258,7 +258,7 @@ def main():
|
|||||||
mcp = FastMCP("my-ai-qwen3-retrieval")
|
mcp = FastMCP("my-ai-qwen3-retrieval")
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
def semantic_search(query: str, top_k: int = 5) -> list[dict]:
|
def semantic_search(query: str, top_k: int = FINAL_RESULTS) -> list[dict]:
|
||||||
"""Search repository concepts, behavior, and implementation patterns first with Qwen3 embedding and reranking. Verify returned passages against current files; use direct reads for known paths."""
|
"""Search repository concepts, behavior, and implementation patterns first with Qwen3 embedding and reranking. Verify returned passages against current files; use direct reads for known paths."""
|
||||||
return index.search(query, top_k)
|
return index.search(query, top_k)
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user