diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 2c9d63882e..1bd172794f 100644 --- a/graphiti_core/search/search_utils.py +++ b/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 diff --git a/tests/utils/search/search_utils_test.py b/tests/utils/search/search_utils_test.py index 6b97daab1f..274b039a26 100644 --- a/tests/utils/search/search_utils_test.py +++ b/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'