Skip to content

Commit

Permalink
fix node distance reranker
Browse files Browse the repository at this point in the history
  • Loading branch information
prasmussen15 committed Dec 6, 2024
1 parent f59a7e6 commit cb0472c
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
8 changes: 6 additions & 2 deletions graphiti_core/search/search_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -631,7 +631,7 @@ async def node_distance_reranker(
) -> list[str]:
# filter out node_uuid center node node uuid
filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
scores: dict[str, float] = {}
scores: dict[str, float] = {center_node_uuid: 0.0}

# Find the shortest path to center node
query = Query("""
Expand All @@ -649,9 +649,13 @@ async def node_distance_reranker(

for result in path_results:
uuid = result['uuid']
score = result['score'] if 'score' in result else float('inf')
score = result['score']
scores[uuid] = score

for uuid in filtered_uuids:
if uuid not in scores:
scores[uuid] = float('inf')

# rerank on shortest distance
filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])

Expand Down
1 change: 1 addition & 0 deletions tests/test_graphiti_int.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ async def test_graphiti_init():
COMBINED_HYBRID_SEARCH_CROSS_ENCODER,
group_ids=['test'],
)

pretty_results = {
'edges': [edge.fact for edge in results.edges],
'nodes': [node.name for node in results.nodes],
Expand Down

0 comments on commit cb0472c

Please sign in to comment.