From 727c499d2c6acd091b8fa20ca99e713c99499215 Mon Sep 17 00:00:00 2001
From: Julian lechner <j.lechner@sobit.at>
Date: Thu, 17 Sep 2026 13:07:47 +0200
Subject: [PATCH] feat(retrieval): rerank semantic search results

---
 scripts/test-retrieval-server.py  | 36 +++++++++++++++++++++++++++++++
 shared/retrieval/requirements.txt |  2 +-
 shared/retrieval/server.py        | 22 ++++++++++++++++---
 3 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 scripts/test-retrieval-server.py

diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py
new file mode 100644
index 0000000..49a1784
--- /dev/null
+++ b/scripts/test-retrieval-server.py
@@ -0,0 +1,36 @@
+import importlib.util
+from pathlib import Path
+import unittest
+
+
+ROOT = Path(__file__).resolve().parent.parent
+SPEC = importlib.util.spec_from_file_location('retrieval_server', ROOT / 'shared/retrieval/server.py')
+MODULE = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(MODULE)
+
+
+class RetrievalServerTests(unittest.TestCase):
+    def test_rerank_uses_cross_encoder_scores(self):
+        index = MODULE.Index.__new__(MODULE.Index)
+
+        class Reranker:
+            def predict(self, pairs):
+                self.pairs = pairs
+                return [0.1, 0.8, 0.5]
+
+        reranker = Reranker()
+        index.reranker = lambda: reranker
+        rows = [
+            {'path': 'a.md', 'score': 0.9, 'text': 'first'},
+            {'path': 'b.md', 'score': 0.8, 'text': 'second'},
+            {'path': 'c.md', 'score': 0.7, 'text': 'third'},
+        ]
+
+        result = index.rerank('query', rows, 2)
+
+        self.assertEqual([item['path'] for item in result], ['b.md', 'c.md'])
+        self.assertEqual(reranker.pairs, [('query', 'first'), ('query', 'second'), ('query', 'third')])
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/shared/retrieval/requirements.txt b/shared/retrieval/requirements.txt
index 95325d9..5f0ec95 100644
--- a/shared/retrieval/requirements.txt
+++ b/shared/retrieval/requirements.txt
@@ -1,2 +1,2 @@
 mcp>=1.9,<2
-sentence-transformers>=4.1,<5
+sentence-transformers>=5.4,<6
diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py
index 7824e6d..3e21da8 100644
--- a/shared/retrieval/server.py
+++ b/shared/retrieval/server.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Local MCP semantic search backed by Qwen3-Embedding-0.6B."""
+"""Local MCP semantic search backed by Qwen3 embedding and reranking."""
 
 from __future__ import annotations
 
@@ -11,6 +11,7 @@ import sqlite3
 
 
 MODEL = "Qwen/Qwen3-Embedding-0.6B"
+RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
 TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
 
 
@@ -31,6 +32,7 @@ class Index:
         self.db.execute("create table if not exists chunks (path text, text text, vector blob)")
         self.db.commit()
         self.model = None
+        self.reranker_model = None
 
     def files(self):
         ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"}
@@ -49,6 +51,19 @@ class Index:
             self.model = SentenceTransformer(MODEL)
         return self.model
 
+    def reranker(self):
+        if self.reranker_model is None:
+            from sentence_transformers import CrossEncoder
+            self.reranker_model = CrossEncoder(RERANKER_MODEL)
+        return self.reranker_model
+
+    def rerank(self, query: str, rows: list[dict], limit: int):
+        candidates = rows[:min(len(rows), max(limit, 20))]
+        scores = self.reranker().predict([(query, item["text"]) for item in candidates])
+        for item, score in zip(candidates, scores):
+            item["score"] = round(float(score), 6)
+        return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit]
+
     def rebuild(self):
         encoder = self.encoder()
         values = []
@@ -73,7 +88,8 @@ class Index:
             import numpy as np
             score = float(np.dot(vector, np.frombuffer(blob, dtype=np.float32)))
             rows.append({"path": path, "score": round(score, 6), "text": text})
-        return sorted(rows, key=lambda item: item["score"], reverse=True)[:max(1, min(top_k, 20))]
+        limit = max(1, min(top_k, 20))
+        return self.rerank(query, sorted(rows, key=lambda item: item["score"], reverse=True), limit)
 
 
 def main():
@@ -91,7 +107,7 @@ def main():
 
     @mcp.tool()
     def semantic_search(query: str, top_k: int = 5) -> list[dict]:
-        """Find repository passages by meaning with Qwen3-Embedding-0.6B."""
+        """Find repository passages with Qwen3 embedding and reranking."""
         return index.search(query, top_k)
 
     @mcp.tool()