graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/search/search.py

@@ -18,10 +18,10 @@ import logging
 from collections import defaultdict
 from time import time
 
-from neo4j import AsyncDriver
-
 from graphiti_core.cross_encoder.client import CrossEncoderClient
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
+from graphiti_core.embedder.client import EMBEDDING_DIM
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import semaphore_gather
@@ -30,6 +30,7 @@ from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
     CommunityReranker,
     CommunitySearchConfig,
+    CommunitySearchMethod,
     EdgeReranker,
     EdgeSearchConfig,
     EdgeSearchMethod,
@@ -73,34 +74,53 @@ async def search(
     center_node_uuid: str | None = None,
     bfs_origin_node_uuids: list[str] | None = None,
     query_vector: list[float] | None = None,
+    driver: GraphDriver | None = None,
 ) -> SearchResults:
     start = time()
 
-    driver = clients.driver
+    driver = driver or clients.driver
     embedder = clients.embedder
     cross_encoder = clients.cross_encoder
 
     if query.strip() == '':
-        return SearchResults(
-            edges=[],
-            nodes=[],
-            episodes=[],
-            communities=[],
+        return SearchResults()
+
+    if (
+        config.edge_config
+        and EdgeSearchMethod.cosine_similarity in config.edge_config.search_methods
+        or config.edge_config
+        and EdgeReranker.mmr == config.edge_config.reranker
+        or config.node_config
+        and NodeSearchMethod.cosine_similarity in config.node_config.search_methods
+        or config.node_config
+        and NodeReranker.mmr == config.node_config.reranker
+        or (
+            config.community_config
+            and CommunitySearchMethod.cosine_similarity in config.community_config.search_methods
         )
-    query_vector = (
-        query_vector
-        if query_vector is not None
-        else await embedder.create(input_data=[query.replace('\n', ' ')])
-    )
+        or (config.community_config and CommunityReranker.mmr == config.community_config.reranker)
+    ):
+        search_vector = (
+            query_vector
+            if query_vector is not None
+            else await embedder.create(input_data=[query.replace('\n', ' ')])
+        )
+    else:
+        search_vector = [0.0] * EMBEDDING_DIM
 
     # if group_ids is empty, set it to None
-    group_ids = group_ids if group_ids else None
-    edges, nodes, episodes, communities = await semaphore_gather(
+    group_ids = group_ids if group_ids and group_ids != [''] else None
+    (
+        (edges, edge_reranker_scores),
+        (nodes, node_reranker_scores),
+        (episodes, episode_reranker_scores),
+        (communities, community_reranker_scores),
+    ) = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.edge_config,
             search_filter,
@@ -113,7 +133,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.node_config,
             search_filter,
@@ -126,7 +146,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.episode_config,
             search_filter,
@@ -137,7 +157,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.community_config,
             config.limit,
@@ -147,9 +167,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
@@ -160,7 +184,7 @@ async def search(
 
 
 async def edge_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -171,51 +195,72 @@ 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(
-            *[
-                edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
-                edge_similarity_search(
-                    driver,
-                    query_vector,
-                    None,
-                    None,
-                    search_filter,
-                    group_ids,
-                    2 * limit,
-                    config.sim_min_score,
-                ),
-                edge_bfs_search(
-                    driver, bfs_origin_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
-                ),
-            ]
+    # Build search tasks based on configured search methods
+    search_tasks = []
+    if EdgeSearchMethod.bm25 in config.search_methods:
+        search_tasks.append(
+            edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
+        )
+    if EdgeSearchMethod.cosine_similarity in config.search_methods:
+        search_tasks.append(
+            edge_similarity_search(
+                driver,
+                query_vector,
+                None,
+                None,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+            )
+        )
+    if EdgeSearchMethod.bfs in config.search_methods:
+        search_tasks.append(
+            edge_bfs_search(
+                driver,
+                bfs_origin_node_uuids,
+                config.bfs_max_depth,
+                search_filter,
+                group_ids,
+                2 * limit,
+            )
         )
-    )
+
+    # Execute only the configured search methods
+    search_results: list[list[EntityEdge]] = []
+    if search_tasks:
+        search_results = list(await semaphore_gather(*search_tasks))
 
     if EdgeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
         source_node_uuids = [edge.source_node_uuid for result in search_results for edge in result]
         search_results.append(
             await edge_bfs_search(
-                driver, source_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
+                driver,
+                source_node_uuids,
+                config.bfs_max_depth,
+                search_filter,
+                group_ids,
+                2 * limit,
             )
         )
 
     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,
@@ -227,12 +272,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,
         )
@@ -245,7 +291,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
         )
 
@@ -257,11 +303,11 @@ 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(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -272,29 +318,54 @@ 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(
-            *[
-                node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
-                node_similarity_search(
-                    driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
-                ),
-                node_bfs_search(
-                    driver, bfs_origin_node_uuids, search_filter, config.bfs_max_depth, 2 * limit
-                ),
-            ]
+    # Build search tasks based on configured search methods
+    search_tasks = []
+    if NodeSearchMethod.bm25 in config.search_methods:
+        search_tasks.append(
+            node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
         )
-    )
+    if NodeSearchMethod.cosine_similarity in config.search_methods:
+        search_tasks.append(
+            node_similarity_search(
+                driver,
+                query_vector,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+            )
+        )
+    if NodeSearchMethod.bfs in config.search_methods:
+        search_tasks.append(
+            node_bfs_search(
+                driver,
+                bfs_origin_node_uuids,
+                search_filter,
+                config.bfs_max_depth,
+                group_ids,
+                2 * limit,
+            )
+        )
+
+    # Execute only the configured search methods
+    search_results: list[list[EntityNode]] = []
+    if search_tasks:
+        search_results = list(await semaphore_gather(*search_tasks))
 
     if NodeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
         origin_node_uuids = [node.uuid for result in search_results for node in result]
         search_results.append(
             await node_bfs_search(
-                driver, origin_node_uuids, search_filter, config.bfs_max_depth, 2 * limit
+                driver,
+                origin_node_uuids,
+                search_filter,
+                config.bfs_max_depth,
+                group_ids,
+                2 * limit,
             )
         )
 
@@ -302,14 +373,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,
@@ -324,27 +396,28 @@ 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(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     _query_vector: list[float],
@@ -353,10 +426,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(
             *[
@@ -369,12 +441,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}
@@ -385,14 +458,15 @@ 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(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -400,9 +474,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(
@@ -421,14 +495,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:
@@ -437,7 +512,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,42 @@ 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)
+
+    @classmethod
+    def merge(cls, results_list: list['SearchResults']) -> 'SearchResults':
+        """
+        Merge multiple SearchResults objects into a single SearchResults object.
+
+        Parameters
+        ----------
+        results_list : list[SearchResults]
+            List of SearchResults objects to merge
+
+        Returns
+        -------
+        SearchResults
+            A single SearchResults object containing all results
+        """
+        if not results_list:
+            return cls()
+
+        merged = cls()
+        for result in results_list:
+            merged.edges.extend(result.edges)
+            merged.edge_reranker_scores.extend(result.edge_reranker_scores)
+            merged.nodes.extend(result.nodes)
+            merged.node_reranker_scores.extend(result.node_reranker_scores)
+            merged.episodes.extend(result.episodes)
+            merged.episode_reranker_scores.extend(result.episode_reranker_scores)
+            merged.communities.extend(result.communities)
+            merged.community_reranker_scores.extend(result.community_reranker_scores)
+
+        return merged