fix(elasticsearch): set size on KNN search to respect top_k (#5910)

This commit is contained in:
Aari
2026-07-31 23:02:27 +08:00
committed by GitHub
parent 760dca6f39
commit 8d45fb3c9a
2 changed files with 43 additions and 1 deletions
+39
View File
@@ -209,6 +209,8 @@ class TestElasticsearchDB(unittest.TestCase):
self.assertEqual(body["knn"]["query_vector"], vectors)
self.assertEqual(body["knn"]["k"], 5)
self.assertEqual(body["knn"]["num_candidates"], 10)
# `size` must be set or ES caps the response at its default of 10 hits.
self.assertEqual(body["size"], 5)
# Verify results
self.assertEqual(len(results), 1)
@@ -216,6 +218,43 @@ class TestElasticsearchDB(unittest.TestCase):
self.assertEqual(results[0].score, 0.8)
self.assertEqual(results[0].payload, {"key1": "value1"})
def test_search_sets_size_to_top_k(self):
"""Regression for #5909: KNN search must set `size=top_k` or Elasticsearch caps
the response at its default of 10 hits regardless of `knn.k`.
Memory._search_vector_store over-fetches a scoring pool of `max(limit * 4, 60)`
candidates for hybrid re-ranking, so even a default search(top_k=20) reaches this
layer with top_k=80. Without `size`, that pool is silently truncated to 10.
"""
self.client_mock.search.return_value = {"hits": {"hits": []}}
# internal_limit in Memory._search_vector_store for a default search(top_k=20)
over_fetch = max(20 * 4, 60)
self.es_db.search(query="", vectors=[[0.1] * 1536], top_k=over_fetch)
body = self.client_mock.search.call_args[1]["body"]
self.assertEqual(body["size"], over_fetch)
self.assertEqual(body["knn"]["k"], over_fetch)
self.assertGreater(body["size"], 10)
def test_search_with_filters_keeps_size(self):
"""Filters must not displace the response size set for KNN search."""
self.client_mock.search.return_value = {"hits": {"hits": []}}
self.es_db.search(
query="",
vectors=[[0.1] * 1536],
top_k=25,
filters={"user_id": "u1"},
)
body = self.client_mock.search.call_args[1]["body"]
self.assertEqual(body["size"], 25)
self.assertEqual(
body["knn"]["filter"],
{"bool": {"must": [{"term": {"metadata.user_id": "u1"}}]}},
)
def test_custom_search_query(self):
# Mock custom search query
self.es_db.custom_search_query = Mock()