graphiti-core 0.17.11__py3-none-any.whl → 0.18.1__py3-none-any.whl

graphiti_core/search/search.py

@@ -80,12 +80,7 @@ async def search(
     cross_encoder = clients.cross_encoder
 
     if query.strip() == '':
-        return SearchResults(
-            edges=[],
-            nodes=[],
-            episodes=[],
-            communities=[],
-        )
+        return SearchResults()
     query_vector = (
         query_vector
         if query_vector is not None
@@ -94,7 +89,12 @@ async def search(
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids and group_ids != [''] else None
-    edges, nodes, episodes, communities = await semaphore_gather(
+    (
+        (edges, edge_reranker_scores),
+        (nodes, node_reranker_scores),
+        (episodes, episode_reranker_scores),
+        (communities, community_reranker_scores),
+    ) = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
@@ -146,9 +146,13 @@ async def search(
 
     results = SearchResults(
         edges=edges,
+        edge_reranker_scores=edge_reranker_scores,
         nodes=nodes,
+        node_reranker_scores=node_reranker_scores,
         episodes=episodes,
+        episode_reranker_scores=episode_reranker_scores,
         communities=communities,
+        community_reranker_scores=community_reranker_scores,
     )
 
     latency = (time() - start) * 1000
@@ -170,9 +174,9 @@ async def edge_search(
     bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EntityEdge]:
+) -> tuple[list[EntityEdge], list[float]]:
     if config is None:
-        return []
+        return [], []
     search_results: list[list[EntityEdge]] = list(
         await semaphore_gather(
             *[
@@ -215,15 +219,16 @@ async def edge_search(
     edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
 
     reranked_uuids: list[str] = []
+    edge_scores: list[float] = []
     if config.reranker == EdgeReranker.rrf or config.reranker == EdgeReranker.episode_mentions:
         search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
 
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, edge_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == EdgeReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_edges(
             driver, list(edge_uuid_map.values())
         )
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, edge_scores = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
@@ -235,12 +240,13 @@ async def edge_search(
         reranked_uuids = [
             fact_to_uuid_map[fact] for fact, score in reranked_facts if score >= reranker_min_score
         ]
+        edge_scores = [score for _, score in reranked_facts if score >= reranker_min_score]
     elif config.reranker == EdgeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
 
         # use rrf as a preliminary sort
-        sorted_result_uuids = rrf(
+        sorted_result_uuids, node_scores = rrf(
             [[edge.uuid for edge in result] for result in search_results],
             min_score=reranker_min_score,
         )
@@ -253,7 +259,7 @@ async def edge_search(
 
         source_uuids = [source_node_uuid for source_node_uuid in source_to_edge_uuid_map]
 
-        reranked_node_uuids = await node_distance_reranker(
+        reranked_node_uuids, edge_scores = await node_distance_reranker(
             driver, source_uuids, center_node_uuid, min_score=reranker_min_score
         )
 
@@ -265,7 +271,7 @@ async def edge_search(
     if config.reranker == EdgeReranker.episode_mentions:
         reranked_edges.sort(reverse=True, key=lambda edge: len(edge.episodes))
 
-    return reranked_edges[:limit]
+    return reranked_edges[:limit], edge_scores[:limit]
 
 
 async def node_search(
@@ -280,9 +286,9 @@ async def node_search(
     bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EntityNode]:
+) -> tuple[list[EntityNode], list[float]]:
     if config is None:
-        return []
+        return [], []
     search_results: list[list[EntityNode]] = list(
         await semaphore_gather(
             *[
@@ -319,14 +325,15 @@ async def node_search(
     node_uuid_map = {node.uuid: node for result in search_results for node in result}
 
     reranked_uuids: list[str] = []
+    node_scores: list[float] = []
     if config.reranker == NodeReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, node_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == NodeReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_nodes(
             driver, list(node_uuid_map.values())
         )
 
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, node_scores = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
@@ -341,23 +348,24 @@ async def node_search(
             for name, score in reranked_node_names
             if score >= reranker_min_score
         ]
+        node_scores = [score for _, score in reranked_node_names if score >= reranker_min_score]
     elif config.reranker == NodeReranker.episode_mentions:
-        reranked_uuids = await episode_mentions_reranker(
+        reranked_uuids, node_scores = await episode_mentions_reranker(
             driver, search_result_uuids, min_score=reranker_min_score
         )
     elif config.reranker == NodeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
-        reranked_uuids = await node_distance_reranker(
+        reranked_uuids, node_scores = await node_distance_reranker(
             driver,
-            rrf(search_result_uuids, min_score=reranker_min_score),
+            rrf(search_result_uuids, min_score=reranker_min_score)[0],
             center_node_uuid,
             min_score=reranker_min_score,
         )
 
     reranked_nodes = [node_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_nodes[:limit]
+    return reranked_nodes[:limit], node_scores[:limit]
 
 
 async def episode_search(
@@ -370,9 +378,9 @@ async def episode_search(
     search_filter: SearchFilters,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EpisodicNode]:
+) -> tuple[list[EpisodicNode], list[float]]:
     if config is None:
-        return []
+        return [], []
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
@@ -385,12 +393,13 @@ async def episode_search(
     episode_uuid_map = {episode.uuid: episode for result in search_results for episode in result}
 
     reranked_uuids: list[str] = []
+    episode_scores: list[float] = []
     if config.reranker == EpisodeReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, episode_scores = rrf(search_result_uuids, min_score=reranker_min_score)
 
     elif config.reranker == EpisodeReranker.cross_encoder:
         # use rrf as a preliminary reranker
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        rrf_result_uuids, episode_scores = rrf(search_result_uuids, min_score=reranker_min_score)
         rrf_results = [episode_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
 
         content_to_uuid_map = {episode.content: episode.uuid for episode in rrf_results}
@@ -401,10 +410,11 @@ async def episode_search(
             for content, score in reranked_contents
             if score >= reranker_min_score
         ]
+        episode_scores = [score for _, score in reranked_contents if score >= reranker_min_score]
 
     reranked_episodes = [episode_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_episodes[:limit]
+    return reranked_episodes[:limit], episode_scores[:limit]
 
 
 async def community_search(
@@ -416,9 +426,9 @@ async def community_search(
     config: CommunitySearchConfig | None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[CommunityNode]:
+) -> tuple[list[CommunityNode], list[float]]:
     if config is None:
-        return []
+        return [], []
 
     search_results: list[list[CommunityNode]] = list(
         await semaphore_gather(
@@ -437,14 +447,15 @@ async def community_search(
     }
 
     reranked_uuids: list[str] = []
+    community_scores: list[float] = []
     if config.reranker == CommunityReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, community_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == CommunityReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_communities(
             driver, list(community_uuid_map.values())
         )
 
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, community_scores = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda, reranker_min_score
         )
     elif config.reranker == CommunityReranker.cross_encoder:
@@ -453,7 +464,8 @@ async def community_search(
         reranked_uuids = [
             name_to_uuid_map[name] for name, score in reranked_nodes if score >= reranker_min_score
         ]
+        community_scores = [score for _, score in reranked_nodes if score >= reranker_min_score]
 
     reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_communities[:limit]
+    return reranked_communities[:limit], community_scores[:limit]

graphiti_core/search/search_config.py

@@ -119,7 +119,11 @@ class SearchConfig(BaseModel):
 
 
 class SearchResults(BaseModel):
-    edges: list[EntityEdge]
-    nodes: list[EntityNode]
-    episodes: list[EpisodicNode]
-    communities: list[CommunityNode]
+    edges: list[EntityEdge] = Field(default_factory=list)
+    edge_reranker_scores: list[float] = Field(default_factory=list)
+    nodes: list[EntityNode] = Field(default_factory=list)
+    node_reranker_scores: list[float] = Field(default_factory=list)
+    episodes: list[EpisodicNode] = Field(default_factory=list)
+    episode_reranker_scores: list[float] = Field(default_factory=list)
+    communities: list[CommunityNode] = Field(default_factory=list)
+    community_reranker_scores: list[float] = Field(default_factory=list)

graphiti_core/search/search_filters.py

@@ -72,7 +72,7 @@ def edge_search_filter_query_constructor(
 
     if filters.edge_types is not None:
         edge_types = filters.edge_types
-        edge_types_filter = '\nAND r.name in $edge_types'
+        edge_types_filter = '\nAND e.name in $edge_types'
         filter_query += edge_types_filter
         filter_params['edge_types'] = edge_types
 
@@ -88,7 +88,7 @@ def edge_search_filter_query_constructor(
                 filter_params['valid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
+                '(e.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -113,7 +113,7 @@ def edge_search_filter_query_constructor(
                 filter_params['invalid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
+                '(e.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -138,7 +138,7 @@ def edge_search_filter_query_constructor(
                 filter_params['created_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
+                '(e.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -163,7 +163,7 @@ def edge_search_filter_query_constructor(
                 filter_params['expired_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
+                '(e.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''