From 2fc4ffcc409345635e15aac5bb5d1eca452fe79e Mon Sep 17 00:00:00 2001 From: Yasyf Mohamedali Date: Wed, 22 Jul 2026 04:20:30 -0700 Subject: [PATCH] Expose semantic cosine scores in search results --- src/semble/search.py | 9 +++++++-- src/semble/types.py | 7 ++++++- src/semble/utils.py | 1 + tests/test_mcp.py | 7 ++++++- tests/test_search.py | 46 +++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 65 insertions(+), 5 deletions(-) diff --git a/src/semble/search.py b/src/semble/search.py index 8d901d699..c540d4338 100644 --- a/src/semble/search.py +++ b/src/semble/search.py @@ -32,7 +32,10 @@ def _search_semantic( query_embedding = model.encode([query]) indices, scores = semantic_index.query(query_embedding, k=top_k, selector=selector)[0] # Vicinity returns cosine distance; convert to similarity so higher = better. - return [SearchResult(chunk=chunks[index], score=1.0 - float(distance)) for index, distance in zip(indices, scores)] + return [ + SearchResult(chunk=chunks[index], score=(similarity := 1.0 - float(distance)), semantic_score=similarity) + for index, distance in zip(indices, scores, strict=True) + ] def _sort_top_k(arr: npt.NDArray, top_k: int) -> npt.NDArray[np.int_]: @@ -129,4 +132,6 @@ def search( else: sorted_by_score = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True) ranked = sorted_by_score[:top_k] - return [SearchResult(chunk=chunk, score=score) for chunk, score in ranked] + return [ + SearchResult(chunk=chunk, score=score, semantic_score=semantic_scores.get(chunk)) for chunk, score in ranked + ] diff --git a/src/semble/types.py b/src/semble/types.py index 97c5551ee..a9988fb59 100644 --- a/src/semble/types.py +++ b/src/semble/types.py @@ -55,10 +55,15 @@ def from_dict(cls: type[Chunk], data: dict[str, Any]) -> Chunk: @dataclass(frozen=True, slots=True) class SearchResult: - """A single search result with score and source.""" + """A single search result. + + ``score`` is the fused or reranked RRF score. ``semantic_score`` is the raw + cosine similarity from the dense leg, or ``None`` for BM25-only results. + """ chunk: Chunk score: float + semantic_score: float | None = None @dataclass(frozen=True, slots=True) diff --git a/src/semble/utils.py b/src/semble/utils.py index 19d786c95..a67d4e3b5 100644 --- a/src/semble/utils.py +++ b/src/semble/utils.py @@ -46,6 +46,7 @@ def format_results(query: str, results: list[SearchResult], max_snippet_lines: i "start_line": r.chunk.start_line, "end_line": r.chunk.end_line, "score": r.score, + "semantic_score": r.semantic_score, } if max_snippet_lines is None: entry["content"] = r.chunk.content diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 77a46046e..9e0e6db88 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -101,9 +101,14 @@ def test_format_results(max_snippet_lines: int | None, has_content: bool, conten assert empty_out == {"query": "query", "results": []} chunks = [make_chunk(f"line1\nline2\nline3\nline4\ndef fn_{i}(): pass", f"f{i}.py") for i in range(3)] - results = [SearchResult(chunk=c, score=round(0.1 * (i + 1), 3)) for i, c in enumerate(chunks)] + semantic_scores = [0.25, None, -0.1] + results = [ + SearchResult(chunk=c, score=round(0.1 * (i + 1), 3), semantic_score=semantic_score) + for i, (c, semantic_score) in enumerate(zip(chunks, semantic_scores)) + ] out = format_results("foo", results, max_snippet_lines) assert out["query"] == "foo" + assert [entry["semantic_score"] for entry in out["results"]] == semantic_scores for entry in out["results"]: assert "file_path" in entry assert "start_line" in entry diff --git a/tests/test_search.py b/tests/test_search.py index fa3d3bd14..09b089ef0 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -11,7 +11,7 @@ from semble.index.dense import SelectableBasicBackend, embed_chunks, load_model from semble.search import _search_bm25, _search_semantic, _sort_top_k, search from semble.tokens import tokenize -from semble.types import Chunk +from semble.types import Chunk, SearchResult from tests.conftest import make_chunk @@ -80,6 +80,23 @@ def test_semantic_search(semantic: SelectableBasicBackend, chunks: list[Chunk], results = _search_semantic("login", mock_model, semantic, chunks, top_k=3, selector=None) assert len(results) > 0 assert all(-1.0 <= r.score <= 1.0 for r in results) + assert all(r.semantic_score == r.score for r in results) + + +def test_semantic_search_converts_cosine_distances_to_scores(chunks: list[Chunk], mock_model: Any) -> None: + """Semantic search converts cosine distances to raw cosine similarities.""" + semantic = MagicMock(spec=SelectableBasicBackend) + semantic.query.return_value = [ + ( + np.array([0, 1, 2], dtype=np.int_), + np.array([0.0, 1.0, 0.5], dtype=np.float32), + ) + ] + + results = _search_semantic("login", mock_model, semantic, chunks, top_k=3, selector=None) + + assert [result.semantic_score for result in results] == [1.0, 0.0, 0.5] + assert [result.score for result in results] == [1.0, 0.0, 0.5] def test_search_hybrid(chunks: list[Chunk], semantic: SelectableBasicBackend, bm25: BM25, mock_model: Any) -> None: @@ -105,6 +122,33 @@ def test_search_hybrid(chunks: list[Chunk], semantic: SelectableBasicBackend, bm assert "module_b.py" in result_locations +@pytest.mark.parametrize("rerank", [True, False]) +def test_search_preserves_semantic_scores(rerank: bool) -> None: + """Hybrid search preserves cosine scores for semantic candidates through both ranking paths.""" + semantic_chunk = make_chunk("def semantic_match(): pass", "semantic.py") + bm25_chunk = make_chunk("def keyword_match(): pass", "keyword.py") + semantic_results = [SearchResult(chunk=semantic_chunk, score=0.8, semantic_score=0.8)] + bm25_results = [SearchResult(chunk=bm25_chunk, score=2.0)] + + with ( + patch("semble.search._search_semantic", return_value=semantic_results), + patch("semble.search._search_bm25", return_value=bm25_results), + ): + results = search( + "unmatched query", + MagicMock(), + MagicMock(), + MagicMock(), + [semantic_chunk, bm25_chunk], + top_k=2, + alpha=0.5, + rerank=rerank, + ) + + scores_by_chunk = {result.chunk: result.semantic_score for result in results} + assert scores_by_chunk == {semantic_chunk: 0.8, bm25_chunk: None} + + @pytest.mark.parametrize( ("search_fn", "query", "top_k"), [