fix(search): rank episode_mentions_reranker by descending mention count (#1342)

graphiti_core/search/search_utils.py

@@ -1872,7 +1872,7 @@ async def episode_mentions_reranker(
     sorted_uuids, _ = rrf(node_uuids)
     scores: dict[str, float] = {}
 
-    # Find the shortest path to center node
+    # Count how many episodes mention each entity
     results, _, _ = await driver.execute_query(
         """
         UNWIND $node_uuids AS node_uuid
@@ -1888,10 +1888,10 @@ async def episode_mentions_reranker(
 
     for uuid in sorted_uuids:
         if uuid not in scores:
-            scores[uuid] = float('inf')
+            scores[uuid] = 0
 
-    # rerank on shortest distance
-    sorted_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
+    # rank by descending mention count (most-mentioned first)
+    sorted_uuids.sort(reverse=True, key=lambda cur_uuid: scores[cur_uuid])
 
     return [uuid for uuid in sorted_uuids if scores[uuid] >= min_score], [
         scores[uuid] for uuid in sorted_uuids if scores[uuid] >= min_score

tests/utils/search/search_utils_test.py

@@ -161,3 +161,49 @@ async def test_hybrid_node_search_with_limit_and_duplicates():
         mock_similarity_search.assert_called_with(
             mock_driver, [0.1, 0.2, 0.3], SearchFilters(), ['1'], 4
         )
+
+
+@pytest.mark.asyncio
+async def test_episode_mentions_reranker_orders_by_descending_mentions():
+    """Nodes mentioned in more episodes must rank higher (regression for #1342)."""
+    from unittest.mock import AsyncMock
+
+    from graphiti_core.search.search_utils import episode_mentions_reranker
+
+    mock_driver = AsyncMock()
+    mock_driver.search_interface = None
+    mock_driver.execute_query.return_value = (
+        [{'uuid': 'few', 'score': 1}, {'uuid': 'many', 'score': 5}],
+        None,
+        None,
+    )
+
+    ranked, scores = await episode_mentions_reranker(
+        mock_driver, node_uuids=[['few', 'many'], ['many', 'few']]
+    )
+
+    assert ranked == ['many', 'few']
+    assert scores == [5, 1]
+
+
+@pytest.mark.asyncio
+async def test_episode_mentions_reranker_unmentioned_nodes_rank_last():
+    """A node with no MENTIONS rows scores 0 and ranks last, not first (regression for #1342)."""
+    from unittest.mock import AsyncMock
+
+    from graphiti_core.search.search_utils import episode_mentions_reranker
+
+    mock_driver = AsyncMock()
+    mock_driver.search_interface = None
+    mock_driver.execute_query.return_value = (
+        [{'uuid': 'mentioned', 'score': 3}],
+        None,
+        None,
+    )
+
+    ranked, _ = await episode_mentions_reranker(
+        mock_driver, node_uuids=[['mentioned', 'unmentioned']]
+    )
+
+    assert ranked[0] == 'mentioned'
+    assert ranked[-1] == 'unmentioned'