From 8d45fb3c9aed8152a07d14e2ee1e5ebf0be05f96 Mon Sep 17 00:00:00 2001 From: Aari Date: Fri, 31 Jul 2026 23:02:27 +0800 Subject: [PATCH] fix(elasticsearch): set size on KNN search to respect top_k (#5910) --- mem0/vector_stores/elasticsearch.py | 5 ++- tests/vector_stores/test_elasticsearch.py | 39 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/mem0/vector_stores/elasticsearch.py b/mem0/vector_stores/elasticsearch.py index 96e520236..6578e9863 100644 --- a/mem0/vector_stores/elasticsearch.py +++ b/mem0/vector_stores/elasticsearch.py @@ -163,7 +163,10 @@ class ElasticsearchDB(VectorStoreBase): search_query = self.custom_search_query(vectors, top_k, filters) else: search_query = { - "knn": {"field": "vector", "query_vector": vectors, "k": top_k, "num_candidates": top_k * 2} + # Without `size`, Elasticsearch caps the response at its default of 10 hits + # regardless of `knn.k`, silently truncating results when top_k > 10. + "size": top_k, + "knn": {"field": "vector", "query_vector": vectors, "k": top_k, "num_candidates": top_k * 2}, } if filters: filter_conditions = [] diff --git a/tests/vector_stores/test_elasticsearch.py b/tests/vector_stores/test_elasticsearch.py index a537ccd6a..57bc9d81f 100644 --- a/tests/vector_stores/test_elasticsearch.py +++ b/tests/vector_stores/test_elasticsearch.py @@ -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()