graphiti-core 0.11.6rc7__py3-none-any.whl → 0.12.0__py3-none-any.whl

graphiti_core/search/search.py

@@ -18,9 +18,8 @@ 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.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
@@ -50,6 +49,9 @@ from graphiti_core.search.search_utils import (
     edge_similarity_search,
     episode_fulltext_search,
     episode_mentions_reranker,
+    get_embeddings_for_communities,
+    get_embeddings_for_edges,
+    get_embeddings_for_nodes,
     maximal_marginal_relevance,
     node_bfs_search,
     node_distance_reranker,
@@ -91,7 +93,7 @@ async def search(
     )
 
     # if group_ids is empty, set it to None
-    group_ids = group_ids if group_ids else None
+    group_ids = group_ids if group_ids and group_ids != [''] else None
     edges, nodes, episodes, communities = await semaphore_gather(
         edge_search(
             driver,
@@ -157,7 +159,7 @@ async def search(
 
 
 async def edge_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -171,7 +173,6 @@ async def edge_search(
 ) -> list[EntityEdge]:
     if config is None:
         return []
-
     search_results: list[list[EntityEdge]] = list(
         await semaphore_gather(
             *[
@@ -209,26 +210,17 @@ async def edge_search(
 
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == EdgeReranker.mmr:
-        await semaphore_gather(
-            *[edge.load_fact_embedding(driver) for result in search_results for edge in result]
+        search_result_uuids_and_vectors = await get_embeddings_for_edges(
+            driver, list(edge_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (edge.uuid, edge.fact_embedding if edge.fact_embedding is not None else [0.0] * 1024)
-            for result in search_results
-            for edge in result
-        ]
         reranked_uuids = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
+            reranker_min_score,
         )
     elif config.reranker == EdgeReranker.cross_encoder:
-        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
-
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
-        rrf_edges = [edge_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
-
-        fact_to_uuid_map = {edge.fact: edge.uuid for edge in rrf_edges}
+        fact_to_uuid_map = {edge.fact: edge.uuid for edge in list(edge_uuid_map.values())[:limit]}
         reranked_facts = await cross_encoder.rank(query, list(fact_to_uuid_map.keys()))
         reranked_uuids = [
             fact_to_uuid_map[fact] for fact, score in reranked_facts if score >= reranker_min_score
@@ -267,7 +259,7 @@ async def edge_search(
 
 
 async def node_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -281,7 +273,6 @@ async def node_search(
 ) -> list[EntityNode]:
     if config is None:
         return []
-
     search_results: list[list[EntityNode]] = list(
         await semaphore_gather(
             *[
@@ -311,30 +302,23 @@ async def node_search(
     if config.reranker == NodeReranker.rrf:
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == NodeReranker.mmr:
-        await semaphore_gather(
-            *[node.load_name_embedding(driver) for result in search_results for node in result]
+        search_result_uuids_and_vectors = await get_embeddings_for_nodes(
+            driver, list(node_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (node.uuid, node.name_embedding if node.name_embedding is not None else [0.0] * 1024)
-            for result in search_results
-            for node in result
-        ]
+
         reranked_uuids = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
+            reranker_min_score,
         )
     elif config.reranker == NodeReranker.cross_encoder:
-        # use rrf as a preliminary reranker
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
-        rrf_results = [node_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
-
-        summary_to_uuid_map = {node.summary: node.uuid for node in rrf_results}
+        name_to_uuid_map = {node.name: node.uuid for node in list(node_uuid_map.values())}
 
-        reranked_summaries = await cross_encoder.rank(query, list(summary_to_uuid_map.keys()))
+        reranked_node_names = await cross_encoder.rank(query, list(name_to_uuid_map.keys()))
         reranked_uuids = [
-            summary_to_uuid_map[fact]
-            for fact, score in reranked_summaries
+            name_to_uuid_map[name]
+            for name, score in reranked_node_names
             if score >= reranker_min_score
         ]
     elif config.reranker == NodeReranker.episode_mentions:
@@ -357,7 +341,7 @@ async def node_search(
 
 
 async def episode_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     _query_vector: list[float],
@@ -369,7 +353,6 @@ async def episode_search(
 ) -> list[EpisodicNode]:
     if config is None:
         return []
-
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
@@ -405,7 +388,7 @@ async def episode_search(
 
 
 async def community_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -437,25 +420,12 @@ async def community_search(
     if config.reranker == CommunityReranker.rrf:
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == CommunityReranker.mmr:
-        await semaphore_gather(
-            *[
-                community.load_name_embedding(driver)
-                for result in search_results
-                for community in result
-            ]
+        search_result_uuids_and_vectors = await get_embeddings_for_communities(
+            driver, list(community_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (
-                community.uuid,
-                community.name_embedding if community.name_embedding is not None else [0.0] * 1024,
-            )
-            for result in search_results
-            for community in result
-        ]
+
         reranked_uuids = maximal_marginal_relevance(
-            query_vector,
-            search_result_uuids_and_vectors,
-            config.mmr_lambda,
+            query_vector, search_result_uuids_and_vectors, config.mmr_lambda, reranker_min_score
         )
     elif config.reranker == CommunityReranker.cross_encoder:
         name_to_uuid_map = {node.name: node.uuid for result in search_results for node in result}

graphiti_core/search/search_filters.py

@@ -42,6 +42,9 @@ class SearchFilters(BaseModel):
     node_labels: list[str] | None = Field(
         default=None, description='List of node labels to filter on'
     )
+    edge_types: list[str] | None = Field(
+        default=None, description='List of edge types to filter on'
+    )
     valid_at: list[list[DateFilter]] | None = Field(default=None)
     invalid_at: list[list[DateFilter]] | None = Field(default=None)
     created_at: list[list[DateFilter]] | None = Field(default=None)
@@ -68,8 +71,19 @@ def edge_search_filter_query_constructor(
     filter_query: LiteralString = ''
     filter_params: dict[str, Any] = {}
 
+    if filters.edge_types is not None:
+        edge_types = filters.edge_types
+        edge_types_filter = '\nAND r.name in $edge_types'
+        filter_query += edge_types_filter
+        filter_params['edge_types'] = edge_types
+
+    if filters.node_labels is not None:
+        node_labels = '|'.join(filters.node_labels)
+        node_label_filter = '\nAND n:' + node_labels + ' AND m:' + node_labels
+        filter_query += node_label_filter
+
     if filters.valid_at is not None:
-        valid_at_filter = ' AND ('
+        valid_at_filter = '\nAND ('
         for i, or_list in enumerate(filters.valid_at):
             for j, date_filter in enumerate(or_list):
                 filter_params['valid_at_' + str(j)] = date_filter.date
@@ -81,12 +95,12 @@ def edge_search_filter_query_constructor(
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             valid_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.valid_at) - 1:
                 valid_at_filter += ')'
             else:
                 valid_at_filter += ' OR '
@@ -106,12 +120,12 @@ def edge_search_filter_query_constructor(
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             invalid_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.invalid_at) - 1:
                 invalid_at_filter += ')'
             else:
                 invalid_at_filter += ' OR '
@@ -131,12 +145,12 @@ def edge_search_filter_query_constructor(
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             created_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.created_at) - 1:
                 created_at_filter += ')'
             else:
                 created_at_filter += ' OR '
@@ -156,12 +170,12 @@ def edge_search_filter_query_constructor(
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             expired_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.expired_at) - 1:
                 expired_at_filter += ')'
             else:
                 expired_at_filter += ' OR '