perf: add FalkorDB HNSW vector indices and fix O(n) fulltext re-match

graphiti_core/driver/falkordb_driver.py

@@ -65,7 +65,7 @@
 from graphiti_core.driver.operations.next_episode_edge_ops import NextEpisodeEdgeOperations
 from graphiti_core.driver.operations.saga_node_ops import SagaNodeOperations
 from graphiti_core.driver.operations.search_ops import SearchOperations
-from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices
+from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices, get_vector_indices
 from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
 
 logger = logging.getLogger(__name__)
@@ -292,14 +292,31 @@ async def delete_all_indexes(self) -> None:
                                 f'DROP FULLTEXT INDEX FOR ()-[e:{label}]-() ON (e.{field_name})'
                             )
                         )
+                elif 'VECTOR' in index_type:
+                    if entity_type == 'NODE':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP VECTOR INDEX FOR (n:{label}) ON (n.{field_name})'
+                            )
+                        )
+                    elif entity_type == 'RELATIONSHIP':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP VECTOR INDEX FOR ()-[e:{label}]-() ON (e.{field_name})'
+                            )
+                        )
 
         if drop_tasks:
             await asyncio.gather(*drop_tasks)
 
     async def build_indices_and_constraints(self, delete_existing=False):
         if delete_existing:
             await self.delete_all_indexes()
-        index_queries = get_range_indices(self.provider) + get_fulltext_indices(self.provider)
+        index_queries = (
+            get_range_indices(self.provider)
+            + get_fulltext_indices(self.provider)
+            + get_vector_indices(self.provider)
+        )
         for query in index_queries:
             await self.execute_query(query)
 

graphiti_core/graph_queries.py

@@ -5,6 +5,8 @@
 supporting index creation, fulltext search, and bulk operations.
 """
 
+import os
+
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import GraphProvider
@@ -140,6 +142,29 @@ def get_fulltext_indices(provider: GraphProvider) -> list[LiteralString]:
     ]
 
 
+def get_vector_indices(provider: GraphProvider) -> list[LiteralString]:
+    """Return HNSW vector index creation queries for the given provider."""
+    if provider == GraphProvider.FALKORDB:
+        dim = int(os.getenv('EMBEDDING_DIM', 1024))
+        # FalkorDB requires dimension at index creation time.
+        # cast() keeps pyright happy with LiteralString expectations.
+        from typing import cast
+
+        return cast(
+            list[LiteralString],
+            [
+                f'CREATE VECTOR INDEX FOR (n:Entity) ON (n.name_embedding)'
+                f" OPTIONS {{dimension: {dim}, similarityFunction: 'cosine'}}",
+                f'CREATE VECTOR INDEX FOR (n:Community) ON (n.name_embedding)'
+                f" OPTIONS {{dimension: {dim}, similarityFunction: 'cosine'}}",
+                f'CREATE VECTOR INDEX FOR ()-[e:RELATES_TO]-() ON (e.fact_embedding)'
+                f" OPTIONS {{dimension: {dim}, similarityFunction: 'cosine'}}",
+            ],
+        )
+
+    return []
+
+
 def get_nodes_query(name: str, query: str, limit: int, provider: GraphProvider) -> str:
     if provider == GraphProvider.FALKORDB:
         label = NEO4J_TO_FALKORDB_MAPPING[name]

graphiti_core/search/search_utils.py

@@ -231,8 +231,8 @@ async def edge_fulltext_search(
                 """
                                 UNWIND $ids as id
                                 MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-                                WHERE e.group_id IN $group_ids 
-                                AND id(e)=id 
+                                WHERE e.group_id IN $group_ids
+                                AND id(e)=id
                                 """
                 + filter_query
                 + """
@@ -265,6 +265,34 @@ async def edge_fulltext_search(
             )
         else:
             return []
+    elif driver.provider == GraphProvider.FALKORDB:
+        # FalkorDB's queryRelationships returns the actual relationship object,
+        # so use startNode/endNode directly instead of re-matching by uuid (which
+        # causes an O(n) scan of all RELATES_TO edges).
+        query = (
+            get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
+            + """
+            YIELD relationship AS e, score
+            WITH e, score, startNode(e) AS n, endNode(e) AS m
+            """
+            + filter_query
+            + """
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider)
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            query=fuzzy_query,
+            limit=limit,
+            routing_='r',
+            **filter_params,
+        )
     else:
         query = (
             get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
@@ -410,6 +438,43 @@ async def edge_similarity_search(
             )
         else:
             return []
+    elif driver.provider == GraphProvider.FALKORDB:
+        # Use HNSW vector index for O(log n) search instead of brute-force scan.
+        # Over-fetch to compensate for post-filtering on group_id, edge_uuids, etc.
+        over_fetch_limit = limit * 10
+
+        post_filter_parts = list(filter_queries)
+        post_filter_parts.append('score > $min_score')
+        post_filter = ' WHERE ' + ' AND '.join(post_filter_parts)
+
+        query = (
+            'CALL db.idx.vector.queryRelationships('
+            "'RELATES_TO', 'fact_embedding', $over_fetch_limit, vecf32($search_vector))"
+            """
+            YIELD relationship AS e, score
+            MATCH (n:Entity)-[e]->(m:Entity)
+            WITH DISTINCT e, n, m, score
+            """
+            + post_filter
+            + """
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider)
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            over_fetch_limit=over_fetch_limit,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
     else:
         query = (
             match_query
@@ -750,6 +815,41 @@ async def node_similarity_search(
             )
         else:
             return []
+    elif driver.provider == GraphProvider.FALKORDB:
+        # Use HNSW vector index for O(log n) search instead of brute-force scan.
+        over_fetch_limit = limit * 10
+
+        post_filter_parts = list(filter_queries)
+        post_filter_parts.append('score > $min_score')
+        post_filter = ' WHERE ' + ' AND '.join(post_filter_parts)
+
+        query = (
+            'CALL db.idx.vector.queryNodes('
+            "'Entity', 'name_embedding', $over_fetch_limit, vecf32($search_vector))"
+            """
+            YIELD node AS n, score
+            WITH n, score
+            """
+            + post_filter
+            + """
+            RETURN
+            """
+            + get_entity_node_return_query(driver.provider)
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            over_fetch_limit=over_fetch_limit,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
     else:
         query = (
             """
@@ -1134,6 +1234,43 @@ async def community_similarity_search(
             )
         else:
             return []
+    elif driver.provider == GraphProvider.FALKORDB:
+        # Use HNSW vector index for O(log n) search instead of brute-force scan.
+        over_fetch_limit = limit * 10
+
+        post_filter_parts: list[str] = []
+        if group_ids is not None:
+            post_filter_parts.append('c.group_id IN $group_ids')
+        post_filter_parts.append('score > $min_score')
+        post_filter = ' WHERE ' + ' AND '.join(post_filter_parts)
+
+        query = (
+            'CALL db.idx.vector.queryNodes('
+            "'Community', 'name_embedding', $over_fetch_limit, vecf32($search_vector))"
+            """
+            YIELD node AS c, score
+            WITH c, score
+            """
+            + post_filter
+            + """
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            over_fetch_limit=over_fetch_limit,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
     else:
         search_vector_var = '$search_vector'
         if driver.provider == GraphProvider.KUZU: