graphiti-core 0.3.15__py3-none-any.whl → 0.3.17__py3-none-any.whl

graphiti_core/search/search.py

@@ -21,6 +21,7 @@ from time import time
 
 from neo4j import AsyncDriver
 
+from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.edges import EntityEdge
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import SearchRerankerError
@@ -39,6 +40,7 @@ from graphiti_core.search.search_config import (
 from graphiti_core.search.search_utils import (
     community_fulltext_search,
     community_similarity_search,
+    edge_bfs_search,
     edge_fulltext_search,
     edge_similarity_search,
     episode_mentions_reranker,
@@ -55,40 +57,49 @@ logger = logging.getLogger(__name__)
 async def search(
     driver: AsyncDriver,
     embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
     group_ids: list[str] | None,
     config: SearchConfig,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
 ) -> SearchResults:
     start = time()
-    query = query.replace('\n', ' ')
+    query_vector = await embedder.create(input=[query.replace('\n', ' ')])
+
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
     edges, nodes, communities = await asyncio.gather(
         edge_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.edge_config,
             center_node_uuid,
+            bfs_origin_node_uuids,
             config.limit,
         ),
         node_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.node_config,
             center_node_uuid,
+            bfs_origin_node_uuids,
             config.limit,
         ),
         community_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.community_config,
+            bfs_origin_node_uuids,
             config.limit,
         ),
     )
@@ -99,27 +110,27 @@ async def search(
         communities=communities,
     )
 
-    end = time()
+    latency = (time() - start) * 1000
 
-    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
+    logger.debug(f'search returned context for query {query} in {latency} ms')
 
     return results
 
 
 async def edge_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: EdgeSearchConfig | None,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[EntityEdge]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[EntityEdge]] = list(
         await asyncio.gather(
             *[
@@ -127,6 +138,7 @@ async def edge_search(
                 edge_similarity_search(
                     driver, query_vector, None, None, group_ids, 2 * limit, config.sim_min_score
                 ),
+                edge_bfs_search(driver, bfs_origin_node_uuids, config.bfs_max_depth),
             ]
         )
     )
@@ -147,6 +159,10 @@ async def edge_search(
         reranked_uuids = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
+    elif config.reranker == EdgeReranker.cross_encoder:
+        fact_to_uuid_map = {edge.fact: edge.uuid for result in search_results for edge in result}
+        reranked_facts = await cross_encoder.rank(query, list(fact_to_uuid_map.keys()))
+        reranked_uuids = [fact_to_uuid_map[fact] for fact, _ in reranked_facts]
     elif config.reranker == EdgeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
@@ -177,18 +193,18 @@ async def edge_search(
 
 async def node_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: NodeSearchConfig | None,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[EntityNode]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[EntityNode]] = list(
         await asyncio.gather(
             *[
@@ -215,6 +231,12 @@ async def node_search(
         reranked_uuids = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
+    elif config.reranker == NodeReranker.cross_encoder:
+        summary_to_uuid_map = {
+            node.summary: node.uuid for result in search_results for node in result
+        }
+        reranked_summaries = await cross_encoder.rank(query, list(summary_to_uuid_map.keys()))
+        reranked_uuids = [summary_to_uuid_map[fact] for fact, _ in reranked_summaries]
     elif config.reranker == NodeReranker.episode_mentions:
         reranked_uuids = await episode_mentions_reranker(driver, search_result_uuids)
     elif config.reranker == NodeReranker.node_distance:
@@ -231,17 +253,17 @@ async def node_search(
 
 async def community_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: CommunitySearchConfig | None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[CommunityNode]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[CommunityNode]] = list(
         await asyncio.gather(
             *[
@@ -273,6 +295,12 @@ async def community_search(
         reranked_uuids = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
+    elif config.reranker == CommunityReranker.cross_encoder:
+        summary_to_uuid_map = {
+            node.summary: node.uuid for result in search_results for node in result
+        }
+        reranked_summaries = await cross_encoder.rank(query, list(summary_to_uuid_map.keys()))
+        reranked_uuids = [summary_to_uuid_map[fact] for fact, _ in reranked_summaries]
 
     reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
 

graphiti_core/search/search_config.py

@@ -20,7 +20,11 @@ from pydantic import BaseModel, Field
 
 from graphiti_core.edges import EntityEdge
 from graphiti_core.nodes import CommunityNode, EntityNode
-from graphiti_core.search.search_utils import DEFAULT_MIN_SCORE, DEFAULT_MMR_LAMBDA
+from graphiti_core.search.search_utils import (
+    DEFAULT_MIN_SCORE,
+    DEFAULT_MMR_LAMBDA,
+    MAX_SEARCH_DEPTH,
+)
 
 DEFAULT_SEARCH_LIMIT = 10
 
@@ -28,11 +32,13 @@ DEFAULT_SEARCH_LIMIT = 10
 class EdgeSearchMethod(Enum):
     cosine_similarity = 'cosine_similarity'
     bm25 = 'bm25'
+    bfs = 'breadth_first_search'
 
 
 class NodeSearchMethod(Enum):
     cosine_similarity = 'cosine_similarity'
     bm25 = 'bm25'
+    bfs = 'breadth_first_search'
 
 
 class CommunitySearchMethod(Enum):
@@ -45,6 +51,7 @@ class EdgeReranker(Enum):
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class NodeReranker(Enum):
@@ -52,11 +59,13 @@ class NodeReranker(Enum):
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class CommunityReranker(Enum):
     rrf = 'reciprocal_rank_fusion'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class EdgeSearchConfig(BaseModel):
@@ -64,6 +73,7 @@ class EdgeSearchConfig(BaseModel):
     reranker: EdgeReranker = Field(default=EdgeReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class NodeSearchConfig(BaseModel):
@@ -71,6 +81,7 @@ class NodeSearchConfig(BaseModel):
     reranker: NodeReranker = Field(default=NodeReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class CommunitySearchConfig(BaseModel):
@@ -78,6 +89,7 @@ class CommunitySearchConfig(BaseModel):
     reranker: CommunityReranker = Field(default=CommunityReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class SearchConfig(BaseModel):

graphiti_core/search/search_config_recipes.py

@@ -48,14 +48,41 @@ COMBINED_HYBRID_SEARCH_MMR = SearchConfig(
     edge_config=EdgeSearchConfig(
         search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
         reranker=EdgeReranker.mmr,
+        mmr_lambda=1,
     ),
     node_config=NodeSearchConfig(
         search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
         reranker=NodeReranker.mmr,
+        mmr_lambda=1,
     ),
     community_config=CommunitySearchConfig(
         search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
         reranker=CommunityReranker.mmr,
+        mmr_lambda=1,
+    ),
+)
+
+# Performs a full-text search, similarity search, and bfs with cross_encoder reranking over edges, nodes, and communities
+COMBINED_HYBRID_SEARCH_CROSS_ENCODER = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[
+            EdgeSearchMethod.bm25,
+            EdgeSearchMethod.cosine_similarity,
+            EdgeSearchMethod.bfs,
+        ],
+        reranker=EdgeReranker.cross_encoder,
+    ),
+    node_config=NodeSearchConfig(
+        search_methods=[
+            NodeSearchMethod.bm25,
+            NodeSearchMethod.cosine_similarity,
+            NodeSearchMethod.bfs,
+        ],
+        reranker=NodeReranker.cross_encoder,
+    ),
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.cross_encoder,
     ),
 )
 
@@ -81,7 +108,6 @@ EDGE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
         search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
         reranker=EdgeReranker.node_distance,
     ),
-    limit=30,
 )
 
 # performs a hybrid search over edges with episode mention reranking