graphiti-core 0.19.0rc3__py3-none-any.whl → 0.20.0__py3-none-any.whl

graphiti_core/search/search_utils.py

@@ -32,15 +32,17 @@ from graphiti_core.graph_queries import (
     get_vector_cosine_func_query,
 )
 from graphiti_core.helpers import (
-    RUNTIME_QUERY,
     lucene_sanitize,
     normalize_l2,
     semaphore_gather,
 )
-from graphiti_core.models.edges.edge_db_queries import ENTITY_EDGE_RETURN
-from graphiti_core.models.nodes.node_db_queries import COMMUNITY_NODE_RETURN, EPISODIC_NODE_RETURN
+from graphiti_core.models.edges.edge_db_queries import get_entity_edge_return_query
+from graphiti_core.models.nodes.node_db_queries import (
+    COMMUNITY_NODE_RETURN,
+    EPISODIC_NODE_RETURN,
+    get_entity_node_return_query,
+)
 from graphiti_core.nodes import (
-    ENTITY_NODE_RETURN,
     CommunityNode,
     EntityNode,
     EpisodicNode,
@@ -78,9 +80,16 @@ def calculate_cosine_similarity(vector1: list[float], vector2: list[float]) -> f
     return dot_product / (norm_vector1 * norm_vector2)
 
 
-def fulltext_query(query: str, group_ids: list[str] | None = None, fulltext_syntax: str = ''):
+def fulltext_query(query: str, group_ids: list[str] | None, driver: GraphDriver):
+    if driver.provider == GraphProvider.KUZU:
+        # Kuzu only supports simple queries.
+        if len(query.split(' ')) > MAX_QUERY_LENGTH:
+            return ''
+        return query
     group_ids_filter_list = (
-        [fulltext_syntax + f'group_id:"{g}"' for g in group_ids] if group_ids is not None else []
+        [driver.fulltext_syntax + f'group_id:"{g}"' for g in group_ids]
+        if group_ids is not None
+        else []
     )
     group_ids_filter = ''
     for f in group_ids_filter_list:
@@ -124,12 +133,12 @@ async def get_mentioned_nodes(
         WHERE episode.uuid IN $uuids
         RETURN DISTINCT
         """
-        + ENTITY_NODE_RETURN,
+        + get_entity_node_return_query(driver.provider),
         uuids=episode_uuids,
         routing_='r',
     )
 
-    nodes = [get_entity_node_from_record(record) for record in records]
+    nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
     return nodes
 
@@ -141,7 +150,7 @@ async def get_communities_by_nodes(
 
     records, _, _ = await driver.execute_query(
         """
-        MATCH (n:Community)-[:HAS_MEMBER]->(m:Entity)
+        MATCH (c:Community)-[:HAS_MEMBER]->(m:Entity)
         WHERE m.uuid IN $uuids
         RETURN DISTINCT
         """
@@ -163,11 +172,32 @@ async def edge_fulltext_search(
     limit=RELEVANT_SCHEMA_LIMIT,
 ) -> list[EntityEdge]:
     # fulltext search over facts
-    fuzzy_query = fulltext_query(query, group_ids, driver.fulltext_syntax)
+    fuzzy_query = fulltext_query(query, group_ids, driver)
+
     if fuzzy_query == '':
         return []
 
-    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
+    match_query = """
+    YIELD relationship AS rel, score
+    MATCH (n:Entity)-[e:RELATES_TO {uuid: rel.uuid}]->(m:Entity)
+    """
+    if driver.provider == GraphProvider.KUZU:
+        match_query = """
+        YIELD node, score
+        MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: node.uuid})-[:RELATES_TO]->(m:Entity)
+        """
+
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
+
+    if group_ids is not None:
+        filter_queries.append('e.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
 
     if driver.provider == GraphProvider.NEPTUNE:
         res = driver.run_aoss_query('edge_name_and_fact', query)  # pyright: ignore reportAttributeAccessIssue
@@ -180,13 +210,14 @@ async def edge_fulltext_search(
             # Match the edge ids and return the values
             query = (
                 """
-                                UNWIND $ids as id
-                                MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-                                WHERE e.group_id IN $group_ids 
-                                AND id(e)=id 
-                                """
+                                        UNWIND $ids as id
+                                        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+                                        WHERE e.group_id IN $group_ids 
+                                        AND id(e)=id 
+                                        """
                 + filter_query
                 + """
+                AND id(e)=id
                 WITH e, id.score as score, startNode(e) AS n, endNode(e) AS m
                 RETURN
                     e.uuid AS uuid,
@@ -208,7 +239,6 @@ async def edge_fulltext_search(
             records, _, _ = await driver.execute_query(
                 query,
                 query=fuzzy_query,
-                group_ids=group_ids,
                 ids=input_ids,
                 limit=limit,
                 routing_='r',
@@ -218,17 +248,14 @@ async def edge_fulltext_search(
             return []
     else:
         query = (
-            get_relationships_query('edge_name_and_fact', provider=driver.provider)
-            + """
-            YIELD relationship AS rel, score
-            MATCH (n:Entity)-[e:RELATES_TO {uuid: rel.uuid}]->(m:Entity)
-            WHERE e.group_id IN $group_ids """
+            get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
+            + match_query
             + filter_query
             + """
             WITH e, score, n, m
             RETURN
             """
-            + ENTITY_EDGE_RETURN
+            + get_entity_edge_return_query(driver.provider)
             + """
             ORDER BY score DESC
             LIMIT $limit
@@ -238,13 +265,12 @@ async def edge_fulltext_search(
         records, _, _ = await driver.execute_query(
             query,
             query=fuzzy_query,
-            group_ids=group_ids,
             limit=limit,
             routing_='r',
             **filter_params,
         )
 
-    edges = [get_entity_edge_from_record(record) for record in records]
+    edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
     return edges
 
@@ -259,32 +285,43 @@ async def edge_similarity_search(
     limit: int = RELEVANT_SCHEMA_LIMIT,
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityEdge]:
-    # vector similarity search over embedded facts
-    query_params: dict[str, Any] = {}
+    match_query = """
+        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+    """
+    if driver.provider == GraphProvider.KUZU:
+        match_query = """
+            MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m:Entity)
+        """
 
-    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
 
-    group_filter_query: LiteralString = 'WHERE e.group_id IS NOT NULL'
     if group_ids is not None:
-        group_filter_query += '\nAND e.group_id IN $group_ids'
-        query_params['group_ids'] = group_ids
+        filter_queries.append('e.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
 
         if source_node_uuid is not None:
-            query_params['source_uuid'] = source_node_uuid
-            group_filter_query += '\nAND (n.uuid = $source_uuid)'
+            filter_params['source_uuid'] = source_node_uuid
+            filter_queries.append('n.uuid = $source_uuid')
 
         if target_node_uuid is not None:
-            query_params['target_uuid'] = target_node_uuid
-            group_filter_query += '\nAND (m.uuid = $target_uuid)'
+            filter_params['target_uuid'] = target_node_uuid
+            filter_queries.append('m.uuid = $target_uuid')
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
+
+    search_vector_var = '$search_vector'
+    if driver.provider == GraphProvider.KUZU:
+        search_vector_var = f'CAST($search_vector AS FLOAT[{len(search_vector)}])'
 
     if driver.provider == GraphProvider.NEPTUNE:
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)            
             """
-            + group_filter_query
+                MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+                """
             + filter_query
             + """
             RETURN DISTINCT id(e) as id, e.fact_embedding as embedding
@@ -296,7 +333,7 @@ async def edge_similarity_search(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
         if len(resp) > 0:
@@ -338,26 +375,22 @@ async def edge_similarity_search(
                 limit=limit,
                 min_score=min_score,
                 routing_='r',
-                **query_params,
+                **filter_params,
             )
         else:
             return []
     else:
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-            """
-            + group_filter_query
+            match_query
             + filter_query
             + """
             WITH DISTINCT e, n, m, """
-            + get_vector_cosine_func_query('e.fact_embedding', '$search_vector', driver.provider)
+            + get_vector_cosine_func_query('e.fact_embedding', search_vector_var, driver.provider)
             + """ AS score
             WHERE score > $min_score
             RETURN
             """
-            + ENTITY_EDGE_RETURN
+            + get_entity_edge_return_query(driver.provider)
             + """
             ORDER BY score DESC
             LIMIT $limit
@@ -370,10 +403,10 @@ async def edge_similarity_search(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
-    edges = [get_entity_edge_from_record(record) for record in records]
+    edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
     return edges
 
@@ -387,70 +420,116 @@ async def edge_bfs_search(
     limit: int = RELEVANT_SCHEMA_LIMIT,
 ) -> list[EntityEdge]:
     # vector similarity search over embedded facts
-    if bfs_origin_node_uuids is None:
+    if bfs_origin_node_uuids is None or len(bfs_origin_node_uuids) == 0:
         return []
 
-    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
 
-    if driver.provider == GraphProvider.NEPTUNE:
-        query = (
+    if group_ids is not None:
+        filter_queries.append('e.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
+
+    if driver.provider == GraphProvider.KUZU:
+        # Kuzu stores entity edges twice with an intermediate node, so we need to match them
+        # separately for the correct BFS depth.
+        depth = bfs_max_depth * 2 - 1
+        match_queries = [
             f"""
-                    UNWIND $bfs_origin_node_uuids AS origin_uuid
-                    MATCH path = (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS *1..{bfs_max_depth}]->(n:Entity)
-                    WHERE origin:Entity OR origin:Episodic
-                    UNWIND relationships(path) AS rel
-                    MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
-                    WHERE e.uuid = rel.uuid
-                """
-            + filter_query
-            + """  
-                        RETURN DISTINCT
-                            e.uuid AS uuid,
-                            e.group_id AS group_id,
-                            startNode(e).uuid AS source_node_uuid,
-                            endNode(e).uuid AS target_node_uuid,
-                            e.created_at AS created_at,
-                            e.name AS name,
-                            e.fact AS fact,
-                            split(e.episodes, ',') AS episodes,
-                            e.expired_at AS expired_at,
-                            e.valid_at AS valid_at,
-                            e.invalid_at AS invalid_at,
-                            properties(e) AS attributes
-                        LIMIT $limit
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH path = (origin:Entity {{uuid: origin_uuid}})-[:RELATES_TO*1..{depth}]->(:RelatesToNode_)
+            UNWIND nodes(path) AS relNode
+            MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {{uuid: relNode.uuid}})-[:RELATES_TO]->(m:Entity)
+            """,
+        ]
+        if bfs_max_depth > 1:
+            depth = (bfs_max_depth - 1) * 2 - 1
+            match_queries.append(f"""
+                UNWIND $bfs_origin_node_uuids AS origin_uuid
+                MATCH path = (origin:Episodic {{uuid: origin_uuid}})-[:MENTIONS]->(:Entity)-[:RELATES_TO*1..{depth}]->(:RelatesToNode_)
+                UNWIND nodes(path) AS relNode
+                MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {{uuid: relNode.uuid}})-[:RELATES_TO]->(m:Entity)
+            """)
+
+        records = []
+        for match_query in match_queries:
+            sub_records, _, _ = await driver.execute_query(
+                match_query
+                + filter_query
+                + """
+                RETURN DISTINCT
                 """
-        )
+                + get_entity_edge_return_query(driver.provider)
+                + """
+                LIMIT $limit
+                """,
+                bfs_origin_node_uuids=bfs_origin_node_uuids,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+            records.extend(sub_records)
     else:
-        query = (
-            f"""
+        if driver.provider == GraphProvider.NEPTUNE:
+            query = (
+                f"""
                 UNWIND $bfs_origin_node_uuids AS origin_uuid
-                MATCH path = (origin:Entity|Episodic {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(:Entity)
+                MATCH path = (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS *1..{bfs_max_depth}]->(n:Entity)
+                WHERE origin:Entity OR origin:Episodic
                 UNWIND relationships(path) AS rel
-                MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
-                WHERE e.uuid = rel.uuid
-                AND e.group_id IN $group_ids
-            """
-            + filter_query
-            + """
-            RETURN DISTINCT
-            """
-            + ENTITY_EDGE_RETURN
-            + """
-            LIMIT $limit
-            """
-        )
+                MATCH (n:Entity)-[e:RELATES_TO {{uuid: rel.uuid}}]-(m:Entity)
+                """
+                + filter_query
+                + """
+                RETURN DISTINCT
+                    e.uuid AS uuid,
+                    e.group_id AS group_id,
+                    startNode(e).uuid AS source_node_uuid,
+                    endNode(e).uuid AS target_node_uuid,
+                    e.created_at AS created_at,
+                    e.name AS name,
+                    e.fact AS fact,
+                    split(e.episodes, ',') AS episodes,
+                    e.expired_at AS expired_at,
+                    e.valid_at AS valid_at,
+                    e.invalid_at AS invalid_at,
+                    properties(e) AS attributes
+                LIMIT $limit
+                """
+            )
+        else:
+            query = (
+                f"""
+                UNWIND $bfs_origin_node_uuids AS origin_uuid
+                MATCH path = (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(:Entity)
+                UNWIND relationships(path) AS rel
+                MATCH (n:Entity)-[e:RELATES_TO {{uuid: rel.uuid}}]-(m:Entity)
+                """
+                + filter_query
+                + """
+                RETURN DISTINCT
+                """
+                + get_entity_edge_return_query(driver.provider)
+                + """
+                LIMIT $limit
+                """
+            )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        bfs_origin_node_uuids=bfs_origin_node_uuids,
-        depth=bfs_max_depth,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
+        records, _, _ = await driver.execute_query(
+            query,
+            bfs_origin_node_uuids=bfs_origin_node_uuids,
+            depth=bfs_max_depth,
+            limit=limit,
+            routing_='r',
+            **filter_params,
+        )
 
-    edges = [get_entity_edge_from_record(record) for record in records]
+    edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
     return edges
 
@@ -461,12 +540,28 @@ async def node_fulltext_search(
     search_filter: SearchFilters,
     group_ids: list[str] | None = None,
     limit=RELEVANT_SCHEMA_LIMIT,
+    use_local_indexes: bool = False,
 ) -> list[EntityNode]:
     # BM25 search to get top nodes
-    fuzzy_query = fulltext_query(query, group_ids, driver.fulltext_syntax)
+    fuzzy_query = fulltext_query(query, group_ids, driver)
     if fuzzy_query == '':
         return []
-    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
+
+    filter_queries, filter_params = node_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
+
+    if group_ids is not None:
+        filter_queries.append('n.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
+
+    yield_query = 'YIELD node AS n, score'
+    if driver.provider == GraphProvider.KUZU:
+        yield_query = 'WITH node AS n, score'
 
     if driver.provider == GraphProvider.NEPTUNE:
         res = driver.run_aoss_query('node_name_and_summary', query, limit=limit)  # pyright: ignore reportAttributeAccessIssue
@@ -479,12 +574,12 @@ async def node_fulltext_search(
             # Match the edge ides and return the values
             query = (
                 """
-                                UNWIND $ids as i
-                                MATCH (n:Entity)
-                                WHERE n.uuid=i.id
-                                RETURN
-                                """
-                + ENTITY_NODE_RETURN
+                                        UNWIND $ids as i
+                                        MATCH (n:Entity)
+                                        WHERE n.uuid=i.id
+                                        RETURN
+                        """
+                + get_entity_node_return_query(driver.provider)
                 + """
                 ORDER BY i.score DESC
                 LIMIT $limit
@@ -494,7 +589,6 @@ async def node_fulltext_search(
                 query,
                 ids=input_ids,
                 query=fuzzy_query,
-                group_ids=group_ids,
                 limit=limit,
                 routing_='r',
                 **filter_params,
@@ -504,36 +598,32 @@ async def node_fulltext_search(
     else:
         index_name = (
             'node_name_and_summary'
-            if not USE_HNSW
+            if not use_local_indexes
             else 'node_name_and_summary_'
             + (group_ids[0].replace('-', '') if group_ids is not None else '')
         )
         query = (
-            get_nodes_query(driver.provider, index_name, '$query')
-            + """
-          YIELD node AS n, score
-          WHERE n:Entity AND n.group_id IN $group_ids
-          """
+            get_nodes_query(index_name, '$query', limit=limit, provider=driver.provider)
+            + yield_query
             + filter_query
             + """
-          WITH n, score
-          ORDER BY score DESC
-          LIMIT $limit
-          RETURN
-          """
-            + ENTITY_NODE_RETURN
+            WITH n, score
+            ORDER BY score DESC
+            LIMIT $limit
+            RETURN
+            """
+            + get_entity_node_return_query(driver.provider)
         )
 
         records, _, _ = await driver.execute_query(
             query,
             query=fuzzy_query,
-            group_ids=group_ids,
             limit=limit,
             routing_='r',
             **filter_params,
         )
 
-    nodes = [get_entity_node_from_record(record) for record in records]
+    nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
     return nodes
 
@@ -545,25 +635,29 @@ async def node_similarity_search(
     group_ids: list[str] | None = None,
     limit=RELEVANT_SCHEMA_LIMIT,
     min_score: float = DEFAULT_MIN_SCORE,
+    use_local_indexes: bool = False,
 ) -> list[EntityNode]:
-    # vector similarity search over entity names
-    query_params: dict[str, Any] = {}
+    filter_queries, filter_params = node_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
 
-    group_filter_query: LiteralString = 'WHERE n.group_id IS NOT NULL'
     if group_ids is not None:
-        group_filter_query += ' AND n.group_id IN $group_ids'
-        query_params['group_ids'] = group_ids
+        filter_queries.append('n.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
 
-    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
+
+    search_vector_var = '$search_vector'
+    if driver.provider == GraphProvider.KUZU:
+        search_vector_var = f'CAST($search_vector AS FLOAT[{len(search_vector)}])'
 
     if driver.provider == GraphProvider.NEPTUNE:
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Entity)
             """
-            + group_filter_query
+                MATCH (n:Entity)
+                """
             + filter_query
             + """
             RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -571,9 +665,8 @@ async def node_similarity_search(
         )
         resp, header, _ = await driver.execute_query(
             query,
-            params=query_params,
+            params=filter_params,
             search_vector=search_vector,
-            group_ids=group_ids,
             limit=limit,
             min_score=min_score,
             routing_='r',
@@ -593,12 +686,12 @@ async def node_similarity_search(
             # Match the edge ides and return the values
             query = (
                 """
-                                    UNWIND $ids as i
-                                    MATCH (n:Entity)
-                                    WHERE id(n)=i.id
-                                    RETURN 
-                                    """
-                + ENTITY_NODE_RETURN
+                                            UNWIND $ids as i
+                                            MATCH (n:Entity)
+                                            WHERE id(n)=i.id
+                                            RETURN 
+                                            """
+                + get_entity_node_return_query(driver.provider)
                 + """
                     ORDER BY i.score DESC
                     LIMIT $limit
@@ -611,11 +704,11 @@ async def node_similarity_search(
                 limit=limit,
                 min_score=min_score,
                 routing_='r',
-                **query_params,
+                **filter_params,
             )
         else:
             return []
-    elif driver.provider == GraphProvider.NEO4J and USE_HNSW:
+    elif driver.provider == GraphProvider.NEO4J and use_local_indexes:
         index_name = 'group_entity_vector_' + (
             group_ids[0].replace('-', '') if group_ids is not None else ''
         )
@@ -623,13 +716,12 @@ async def node_similarity_search(
             f"""
                     CALL db.index.vector.queryNodes('{index_name}', {limit}, $search_vector) YIELD node AS n, score
                     """
-            + group_filter_query
             + filter_query
             + """
                     AND score > $min_score
                     RETURN
                     """
-            + ENTITY_NODE_RETURN
+            + get_entity_node_return_query(driver.provider)
             + """
                     ORDER BY score DESC
                     LIMIT $limit
@@ -642,25 +734,23 @@ async def node_similarity_search(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
     else:
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Entity)
             """
-            + group_filter_query
+                MATCH (n:Entity)
+                """
             + filter_query
             + """
             WITH n, """
-            + get_vector_cosine_func_query('n.name_embedding', '$search_vector', driver.provider)
+            + get_vector_cosine_func_query('n.name_embedding', search_vector_var, driver.provider)
             + """ AS score
             WHERE score > $min_score
             RETURN
             """
-            + ENTITY_NODE_RETURN
+            + get_entity_node_return_query(driver.provider)
             + """
             ORDER BY score DESC
             LIMIT $limit
@@ -673,10 +763,10 @@ async def node_similarity_search(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
-    nodes = [get_entity_node_from_record(record) for record in records]
+    nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
     return nodes
 
@@ -689,56 +779,82 @@ async def node_bfs_search(
     group_ids: list[str] | None = None,
     limit: int = RELEVANT_SCHEMA_LIMIT,
 ) -> list[EntityNode]:
-    # vector similarity search over entity names
-    if bfs_origin_node_uuids is None:
+    if bfs_origin_node_uuids is None or len(bfs_origin_node_uuids) == 0 or bfs_max_depth < 1:
         return []
 
-    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
+    filter_queries, filter_params = node_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
+
+    if group_ids is not None:
+        filter_queries.append('n.group_id IN $group_ids')
+        filter_queries.append('origin.group_id IN $group_ids')
+        filter_params['group_ids'] = group_ids
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' AND ' + (' AND '.join(filter_queries))
+
+    match_queries = [
+        f"""
+        UNWIND $bfs_origin_node_uuids AS origin_uuid
+        MATCH (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(n:Entity)
+        WHERE n.group_id = origin.group_id
+        """
+    ]
 
     if driver.provider == GraphProvider.NEPTUNE:
-        query = (
+        match_queries = [
             f"""
-                UNWIND $bfs_origin_node_uuids AS origin_uuid            
-                MATCH (origin {{uuid: origin_uuid}})-[e:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(n:Entity)
-                WHERE origin:Entity OR origin.Episode
-                AND n.group_id = origin.group_id
-            """
-            + filter_query
-            + """
-            RETURN
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH (origin {{uuid: origin_uuid}})-[e:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(n:Entity)
+            WHERE origin:Entity OR origin.Episode
+            AND n.group_id = origin.group_id
             """
-            + ENTITY_NODE_RETURN
-            + """
-            LIMIT $limit
+        ]
+
+    if driver.provider == GraphProvider.KUZU:
+        depth = bfs_max_depth * 2
+        match_queries = [
             """
-        )
-    else:
-        query = (
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH (origin:Episodic {uuid: origin_uuid})-[:MENTIONS]->(n:Entity)
+            WHERE n.group_id = origin.group_id
+            """,
             f"""
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH (origin:Entity {{uuid: origin_uuid}})-[:RELATES_TO*2..{depth}]->(n:Entity)
+            WHERE n.group_id = origin.group_id
+            """,
+        ]
+        if bfs_max_depth > 1:
+            depth = (bfs_max_depth - 1) * 2
+            match_queries.append(f"""
                 UNWIND $bfs_origin_node_uuids AS origin_uuid
-                MATCH (origin:Entity|Episodic {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{bfs_max_depth}]->(n:Entity)
+                MATCH (origin:Episodic {{uuid: origin_uuid}})-[:MENTIONS]->(:Entity)-[:RELATES_TO*2..{depth}]->(n:Entity)
                 WHERE n.group_id = origin.group_id
-                AND origin.group_id IN $group_ids
-            """
+            """)
+
+    records = []
+    for match_query in match_queries:
+        sub_records, _, _ = await driver.execute_query(
+            match_query
             + filter_query
             + """
             RETURN
             """
-            + ENTITY_NODE_RETURN
+            + get_entity_node_return_query(driver.provider)
             + """
             LIMIT $limit
-            """
+            """,
+            bfs_origin_node_uuids=bfs_origin_node_uuids,
+            limit=limit,
+            routing_='r',
+            **filter_params,
         )
+        records.extend(sub_records)
 
-    records, _, _ = await driver.execute_query(
-        query,
-        bfs_origin_node_uuids=bfs_origin_node_uuids,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
-    nodes = [get_entity_node_from_record(record) for record in records]
+    nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
     return nodes
 
@@ -749,12 +865,19 @@ async def episode_fulltext_search(
     _search_filter: SearchFilters,
     group_ids: list[str] | None = None,
     limit=RELEVANT_SCHEMA_LIMIT,
+    use_local_indexes: bool = False,
 ) -> list[EpisodicNode]:
     # BM25 search to get top episodes
-    fuzzy_query = fulltext_query(query, group_ids, driver.fulltext_syntax)
+    fuzzy_query = fulltext_query(query, group_ids, driver)
     if fuzzy_query == '':
         return []
 
+    filter_params: dict[str, Any] = {}
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += '\nAND e.group_id IN $group_ids'
+        filter_params['group_ids'] = group_ids
+
     if driver.provider == GraphProvider.NEPTUNE:
         res = driver.run_aoss_query('episode_content', query, limit=limit)  # pyright: ignore reportAttributeAccessIssue
         if res['hits']['total']['value'] > 0:
@@ -768,7 +891,7 @@ async def episode_fulltext_search(
                 UNWIND $ids as i
                 MATCH (e:Episodic)
                 WHERE e.uuid=i.id
-            RETURN 
+            RETURN
                     e.content AS content,
                     e.created_at AS created_at,
                     e.valid_at AS valid_at,
@@ -785,26 +908,28 @@ async def episode_fulltext_search(
                 query,
                 ids=input_ids,
                 query=fuzzy_query,
-                group_ids=group_ids,
                 limit=limit,
                 routing_='r',
+                **filter_params,
             )
         else:
             return []
     else:
         index_name = (
             'episode_content'
-            if not USE_HNSW
+            if not use_local_indexes
             else 'episode_content_'
             + (group_ids[0].replace('-', '') if group_ids is not None else '')
         )
         query = (
-            get_nodes_query(driver.provider, index_name, '$query')
+            get_nodes_query(index_name, '$query', limit=limit, provider=driver.provider)
             + """
             YIELD node AS episode, score
             MATCH (e:Episodic)
             WHERE e.uuid = episode.uuid
-            AND e.group_id IN $group_ids
+            """
+            + group_filter_query
+            + """
             RETURN
             """
             + EPISODIC_NODE_RETURN
@@ -815,12 +940,9 @@ async def episode_fulltext_search(
         )
 
         records, _, _ = await driver.execute_query(
-            query,
-            query=fuzzy_query,
-            group_ids=group_ids,
-            limit=limit,
-            routing_='r',
+            query, query=fuzzy_query, limit=limit, routing_='r', **filter_params
         )
+
     episodes = [get_episodic_node_from_record(record) for record in records]
 
     return episodes
@@ -833,10 +955,20 @@ async def community_fulltext_search(
     limit=RELEVANT_SCHEMA_LIMIT,
 ) -> list[CommunityNode]:
     # BM25 search to get top communities
-    fuzzy_query = fulltext_query(query, group_ids, driver.fulltext_syntax)
+    fuzzy_query = fulltext_query(query, group_ids, driver)
     if fuzzy_query == '':
         return []
 
+    filter_params: dict[str, Any] = {}
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query = 'WHERE c.group_id IN $group_ids'
+        filter_params['group_ids'] = group_ids
+
+    yield_query = 'YIELD node AS c, score'
+    if driver.provider == GraphProvider.KUZU:
+        yield_query = 'WITH node AS c, score'
+
     if driver.provider == GraphProvider.NEPTUNE:
         res = driver.run_aoss_query('community_name', query, limit=limit)  # pyright: ignore reportAttributeAccessIssue
         if res['hits']['total']['value'] > 0:
@@ -852,9 +984,9 @@ async def community_fulltext_search(
                 WHERE comm.uuid=i.id
                 RETURN
                     comm.uuid AS uuid,
-                    comm.group_id AS group_id, 
-                    comm.name AS name, 
-                    comm.created_at AS created_at, 
+                    comm.group_id AS group_id,
+                    comm.name AS name,
+                    comm.created_at AS created_at,
                     comm.summary AS summary,
                     [x IN split(comm.name_embedding, ",") | toFloat(x)]AS name_embedding
                 ORDER BY i.score DESC
@@ -864,18 +996,21 @@ async def community_fulltext_search(
                 query,
                 ids=input_ids,
                 query=fuzzy_query,
-                group_ids=group_ids,
                 limit=limit,
                 routing_='r',
+                **filter_params,
             )
         else:
             return []
     else:
         query = (
-            get_nodes_query(driver.provider, 'community_name', '$query')
+            get_nodes_query('community_name', '$query', limit=limit, provider=driver.provider)
+            + yield_query
+            + """
+            WITH c, score
+            """
+            + group_filter_query
             + """
-            YIELD node AS n, score
-            WHERE n.group_id IN $group_ids
             RETURN
             """
             + COMMUNITY_NODE_RETURN
@@ -886,12 +1021,9 @@ async def community_fulltext_search(
         )
 
         records, _, _ = await driver.execute_query(
-            query,
-            query=fuzzy_query,
-            group_ids=group_ids,
-            limit=limit,
-            routing_='r',
+            query, query=fuzzy_query, limit=limit, routing_='r', **filter_params
         )
+
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -909,15 +1041,14 @@ async def community_similarity_search(
 
     group_filter_query: LiteralString = ''
     if group_ids is not None:
-        group_filter_query += 'WHERE n.group_id IN $group_ids'
+        group_filter_query += ' WHERE c.group_id IN $group_ids'
         query_params['group_ids'] = group_ids
 
     if driver.provider == GraphProvider.NEPTUNE:
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Community)
             """
+                MATCH (n:Community)
+                """
             + group_filter_query
             + """
             RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -951,8 +1082,8 @@ async def community_similarity_search(
                     RETURN
                         comm.uuid As uuid,
                         comm.group_id AS group_id,
-                        comm.name AS name, 
-                        comm.created_at AS created_at, 
+                        comm.name AS name,
+                        comm.created_at AS created_at,
                         comm.summary AS summary,
                         comm.name_embedding AS name_embedding
                     ORDER BY i.score DESC
@@ -970,16 +1101,19 @@ async def community_similarity_search(
         else:
             return []
     else:
+        search_vector_var = '$search_vector'
+        if driver.provider == GraphProvider.KUZU:
+            search_vector_var = f'CAST($search_vector AS FLOAT[{len(search_vector)}])'
+
         query = (
-            RUNTIME_QUERY
-            + """
-            MATCH (n:Community)
             """
+                MATCH (c:Community)
+                """
             + group_filter_query
             + """
-            WITH n,
+            WITH c,
             """
-            + get_vector_cosine_func_query('n.name_embedding', '$search_vector', driver.provider)
+            + get_vector_cosine_func_query('c.name_embedding', search_vector_var, driver.provider)
             + """ AS score
             WHERE score > $min_score
             RETURN
@@ -999,6 +1133,7 @@ async def community_similarity_search(
             routing_='r',
             **query_params,
         )
+
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -1089,67 +1224,127 @@ async def get_relevant_nodes(
         return []
 
     group_id = nodes[0].group_id
-
-    # vector similarity search over entity names
-    query_params: dict[str, Any] = {}
-
-    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
-
-    query = (
-        RUNTIME_QUERY
-        + """
-        UNWIND $nodes AS node
-        MATCH (n:Entity {group_id: $group_id})
-        """
-        + filter_query
-        + """
-        WITH node, n, """
-        + get_vector_cosine_func_query('n.name_embedding', 'node.name_embedding', driver.provider)
-        + """ AS score
-        WHERE score > $min_score
-        WITH node, collect(n)[..$limit] AS top_vector_nodes, collect(n.uuid) AS vector_node_uuids
-        """
-        + get_nodes_query(driver.provider, 'node_name_and_summary', 'node.fulltext_query')
-        + """
-        YIELD node AS m
-        WHERE m.group_id = $group_id
-        WITH node, top_vector_nodes, vector_node_uuids, collect(m) AS fulltext_nodes
-
-        WITH node,
-             top_vector_nodes,
-             [m IN fulltext_nodes WHERE NOT m.uuid IN vector_node_uuids] AS filtered_fulltext_nodes
-
-        WITH node, top_vector_nodes + filtered_fulltext_nodes AS combined_nodes
-
-        UNWIND combined_nodes AS combined_node
-        WITH node, collect(DISTINCT combined_node) AS deduped_nodes
-
-        RETURN
-          node.uuid AS search_node_uuid,
-          [x IN deduped_nodes | {
-            uuid: x.uuid,
-            name: x.name,
-            name_embedding: x.name_embedding,
-            group_id: x.group_id,
-            created_at: x.created_at,
-            summary: x.summary,
-            labels: labels(x),
-            attributes: properties(x)
-          }] AS matches
-        """
-    )
-
     query_nodes = [
         {
             'uuid': node.uuid,
             'name': node.name,
             'name_embedding': node.name_embedding,
-            'fulltext_query': fulltext_query(node.name, [node.group_id], driver.fulltext_syntax),
+            'fulltext_query': fulltext_query(node.name, [node.group_id], driver),
         }
         for node in nodes
     ]
 
+    filter_queries, filter_params = node_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
+
+    filter_query = ''
+    if filter_queries:
+        filter_query = 'WHERE ' + (' AND '.join(filter_queries))
+
+    if driver.provider == GraphProvider.KUZU:
+        embedding_size = len(nodes[0].name_embedding) if nodes[0].name_embedding is not None else 0
+        if embedding_size == 0:
+            return []
+
+        # FIXME: Kuzu currently does not support using variables such as `node.fulltext_query` as an input to FTS, which means `get_relevant_nodes()` won't work with Kuzu as the graph driver.
+        query = (
+            """
+                UNWIND $nodes AS node
+                MATCH (n:Entity {group_id: $group_id})
+                """
+            + filter_query
+            + """
+            WITH node, n, """
+            + get_vector_cosine_func_query(
+                'n.name_embedding',
+                f'CAST(node.name_embedding AS FLOAT[{embedding_size}])',
+                driver.provider,
+            )
+            + """ AS score
+            WHERE score > $min_score
+            WITH node, collect(n)[:$limit] AS top_vector_nodes, collect(n.uuid) AS vector_node_uuids
+            """
+            + get_nodes_query(
+                'node_name_and_summary',
+                'node.fulltext_query',
+                limit=limit,
+                provider=driver.provider,
+            )
+            + """
+            WITH node AS m
+            WHERE m.group_id = $group_id AND NOT m.uuid IN vector_node_uuids
+            WITH node, top_vector_nodes, collect(m) AS fulltext_nodes
+
+            WITH node, list_concat(top_vector_nodes, fulltext_nodes) AS combined_nodes
+
+            UNWIND combined_nodes AS x
+            WITH node, collect(DISTINCT {
+                uuid: x.uuid,
+                name: x.name,
+                name_embedding: x.name_embedding,
+                group_id: x.group_id,
+                created_at: x.created_at,
+                summary: x.summary,
+                labels: x.labels,
+                attributes: x.attributes
+            }) AS matches
+
+            RETURN
+            node.uuid AS search_node_uuid, matches
+            """
+        )
+    else:
+        query = (
+            """
+                UNWIND $nodes AS node
+                MATCH (n:Entity {group_id: $group_id})
+                """
+            + filter_query
+            + """
+            WITH node, n, """
+            + get_vector_cosine_func_query(
+                'n.name_embedding', 'node.name_embedding', driver.provider
+            )
+            + """ AS score
+            WHERE score > $min_score
+            WITH node, collect(n)[..$limit] AS top_vector_nodes, collect(n.uuid) AS vector_node_uuids
+            """
+            + get_nodes_query(
+                'node_name_and_summary',
+                'node.fulltext_query',
+                limit=limit,
+                provider=driver.provider,
+            )
+            + """
+            YIELD node AS m
+            WHERE m.group_id = $group_id
+            WITH node, top_vector_nodes, vector_node_uuids, collect(m) AS fulltext_nodes
+
+            WITH node,
+                top_vector_nodes,
+                [m IN fulltext_nodes WHERE NOT m.uuid IN vector_node_uuids] AS filtered_fulltext_nodes
+
+            WITH node, top_vector_nodes + filtered_fulltext_nodes AS combined_nodes
+
+            UNWIND combined_nodes AS combined_node
+            WITH node, collect(DISTINCT combined_node) AS deduped_nodes
+
+            RETURN
+            node.uuid AS search_node_uuid,
+            [x IN deduped_nodes | {
+                uuid: x.uuid,
+                name: x.name,
+                name_embedding: x.name_embedding,
+                group_id: x.group_id,
+                created_at: x.created_at,
+                summary: x.summary,
+                labels: labels(x),
+                attributes: properties(x)
+            }] AS matches
+            """
+        )
+
     results, _, _ = await driver.execute_query(
         query,
         nodes=query_nodes,
@@ -1157,12 +1352,12 @@ async def get_relevant_nodes(
         limit=limit,
         min_score=min_score,
         routing_='r',
-        **query_params,
+        **filter_params,
     )
 
     relevant_nodes_dict: dict[str, list[EntityNode]] = {
         result['search_node_uuid']: [
-            get_entity_node_from_record(record) for record in result['matches']
+            get_entity_node_from_record(record, driver.provider) for record in result['matches']
         ]
         for result in results
     }
@@ -1182,22 +1377,24 @@ async def get_relevant_edges(
     if len(edges) == 0:
         return []
 
-    query_params: dict[str, Any] = {}
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
 
-    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' WHERE ' + (' AND '.join(filter_queries))
 
     if driver.provider == GraphProvider.NEPTUNE:
         query = (
-            RUNTIME_QUERY
-            + """
-            UNWIND $edges AS edge
-            MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
             """
+                UNWIND $edges AS edge
+                MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
+                """
             + filter_query
             + """
             WITH e, edge
-            RETURN DISTINCT id(e) as id, e.fact_embedding as source_embedding, edge.uuid as search_edge_uuid, 
+            RETURN DISTINCT id(e) as id, e.fact_embedding as source_embedding, edge.uuid as search_edge_uuid,
             edge.fact_embedding as target_embedding
             """
         )
@@ -1207,7 +1404,7 @@ async def get_relevant_edges(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
         # Calculate Cosine similarity then return the edge ids
@@ -1220,7 +1417,7 @@ async def get_relevant_edges(
                 input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']})
 
         # Match the edge ides and return the values
-        query = """                      
+        query = """
         UNWIND $ids AS edge
         MATCH ()-[e]->()
         WHERE id(e) = edge.id
@@ -1246,49 +1443,93 @@ async def get_relevant_edges(
 
         results, _, _ = await driver.execute_query(
             query,
-            params=query_params,
             ids=input_ids,
             edges=[edge.model_dump() for edge in edges],
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
     else:
-        query = (
-            RUNTIME_QUERY
-            + """
-            UNWIND $edges AS edge
-            MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
-            """
-            + filter_query
-            + """
-            WITH e, edge, """
-            + get_vector_cosine_func_query(
-                'e.fact_embedding', 'edge.fact_embedding', driver.provider
+        if driver.provider == GraphProvider.KUZU:
+            embedding_size = (
+                len(edges[0].fact_embedding) if edges[0].fact_embedding is not None else 0
+            )
+            if embedding_size == 0:
+                return []
+
+            query = (
+                """
+                    UNWIND $edges AS edge
+                    MATCH (n:Entity {uuid: edge.source_node_uuid})-[:RELATES_TO]-(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]-(m:Entity {uuid: edge.target_node_uuid})
+                    """
+                + filter_query
+                + """
+                WITH e, edge, n, m, """
+                + get_vector_cosine_func_query(
+                    'e.fact_embedding',
+                    f'CAST(edge.fact_embedding AS FLOAT[{embedding_size}])',
+                    driver.provider,
+                )
+                + """ AS score
+                WHERE score > $min_score
+                WITH e, edge, n, m, score
+                ORDER BY score DESC
+                LIMIT $limit
+                RETURN
+                    edge.uuid AS search_edge_uuid,
+                    collect({
+                        uuid: e.uuid,
+                        source_node_uuid: n.uuid,
+                        target_node_uuid: m.uuid,
+                        created_at: e.created_at,
+                        name: e.name,
+                        group_id: e.group_id,
+                        fact: e.fact,
+                        fact_embedding: e.fact_embedding,
+                        episodes: e.episodes,
+                        expired_at: e.expired_at,
+                        valid_at: e.valid_at,
+                        invalid_at: e.invalid_at,
+                        attributes: e.attributes
+                    }) AS matches
+                """
+            )
+        else:
+            query = (
+                """
+                    UNWIND $edges AS edge
+                    MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
+                    """
+                + filter_query
+                + """
+                WITH e, edge, """
+                + get_vector_cosine_func_query(
+                    'e.fact_embedding', 'edge.fact_embedding', driver.provider
+                )
+                + """ AS score
+                WHERE score > $min_score
+                WITH edge, e, score
+                ORDER BY score DESC
+                RETURN
+                    edge.uuid AS search_edge_uuid,
+                    collect({
+                        uuid: e.uuid,
+                        source_node_uuid: startNode(e).uuid,
+                        target_node_uuid: endNode(e).uuid,
+                        created_at: e.created_at,
+                        name: e.name,
+                        group_id: e.group_id,
+                        fact: e.fact,
+                        fact_embedding: e.fact_embedding,
+                        episodes: e.episodes,
+                        expired_at: e.expired_at,
+                        valid_at: e.valid_at,
+                        invalid_at: e.invalid_at,
+                        attributes: properties(e)
+                    })[..$limit] AS matches
+                """
             )
-            + """ AS score
-            WHERE score > $min_score
-            WITH edge, e, score
-            ORDER BY score DESC
-            RETURN edge.uuid AS search_edge_uuid,
-                collect({
-                    uuid: e.uuid,
-                    source_node_uuid: startNode(e).uuid,
-                    target_node_uuid: endNode(e).uuid,
-                    created_at: e.created_at,
-                    name: e.name,
-                    group_id: e.group_id,
-                    fact: e.fact,
-                    fact_embedding: e.fact_embedding,
-                    episodes: e.episodes,
-                    expired_at: e.expired_at,
-                    valid_at: e.valid_at,
-                    invalid_at: e.invalid_at,
-                    attributes: properties(e)
-                })[..$limit] AS matches
-            """
-        )
 
         results, _, _ = await driver.execute_query(
             query,
@@ -1296,12 +1537,12 @@ async def get_relevant_edges(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
     relevant_edges_dict: dict[str, list[EntityEdge]] = {
         result['search_edge_uuid']: [
-            get_entity_edge_from_record(record) for record in result['matches']
+            get_entity_edge_from_record(record, driver.provider) for record in result['matches']
         ]
         for result in results
     }
@@ -1321,19 +1562,21 @@ async def get_edge_invalidation_candidates(
     if len(edges) == 0:
         return []
 
-    query_params: dict[str, Any] = {}
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
+    )
 
-    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
+    filter_query = ''
+    if filter_queries:
+        filter_query = ' AND ' + (' AND '.join(filter_queries))
 
     if driver.provider == GraphProvider.NEPTUNE:
         query = (
-            RUNTIME_QUERY
-            + """
-            UNWIND $edges AS edge
-            MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
-            WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
             """
+                UNWIND $edges AS edge
+                MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
+                WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
+                """
             + filter_query
             + """
             WITH e, edge
@@ -1348,7 +1591,7 @@ async def get_edge_invalidation_candidates(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
 
         # Calculate Cosine similarity then return the edge ids
@@ -1361,7 +1604,7 @@ async def get_edge_invalidation_candidates(
                 input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']})
 
         # Match the edge ides and return the values
-        query = """                      
+        query = """
         UNWIND $ids AS edge
         MATCH ()-[e]->()
         WHERE id(e) = edge.id
@@ -1391,44 +1634,90 @@ async def get_edge_invalidation_candidates(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
     else:
-        query = (
-            RUNTIME_QUERY
-            + """
-            UNWIND $edges AS edge
-            MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
-            WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
-            """
-            + filter_query
-            + """
-            WITH edge, e, """
-            + get_vector_cosine_func_query(
-                'e.fact_embedding', 'edge.fact_embedding', driver.provider
+        if driver.provider == GraphProvider.KUZU:
+            embedding_size = (
+                len(edges[0].fact_embedding) if edges[0].fact_embedding is not None else 0
+            )
+            if embedding_size == 0:
+                return []
+
+            query = (
+                """
+                    UNWIND $edges AS edge
+                    MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity)
+                    WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid])
+                    """
+                + filter_query
+                + """
+                WITH edge, e, n, m, """
+                + get_vector_cosine_func_query(
+                    'e.fact_embedding',
+                    f'CAST(edge.fact_embedding AS FLOAT[{embedding_size}])',
+                    driver.provider,
+                )
+                + """ AS score
+                WHERE score > $min_score
+                WITH edge, e, n, m, score
+                ORDER BY score DESC
+                LIMIT $limit
+                RETURN
+                    edge.uuid AS search_edge_uuid,
+                    collect({
+                        uuid: e.uuid,
+                        source_node_uuid: n.uuid,
+                        target_node_uuid: m.uuid,
+                        created_at: e.created_at,
+                        name: e.name,
+                        group_id: e.group_id,
+                        fact: e.fact,
+                        fact_embedding: e.fact_embedding,
+                        episodes: e.episodes,
+                        expired_at: e.expired_at,
+                        valid_at: e.valid_at,
+                        invalid_at: e.invalid_at,
+                        attributes: e.attributes
+                    }) AS matches
+                """
+            )
+        else:
+            query = (
+                """
+                    UNWIND $edges AS edge
+                    MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
+                    WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
+                    """
+                + filter_query
+                + """
+                WITH edge, e, """
+                + get_vector_cosine_func_query(
+                    'e.fact_embedding', 'edge.fact_embedding', driver.provider
+                )
+                + """ AS score
+                WHERE score > $min_score
+                WITH edge, e, score
+                ORDER BY score DESC
+                RETURN
+                    edge.uuid AS search_edge_uuid,
+                    collect({
+                        uuid: e.uuid,
+                        source_node_uuid: startNode(e).uuid,
+                        target_node_uuid: endNode(e).uuid,
+                        created_at: e.created_at,
+                        name: e.name,
+                        group_id: e.group_id,
+                        fact: e.fact,
+                        fact_embedding: e.fact_embedding,
+                        episodes: e.episodes,
+                        expired_at: e.expired_at,
+                        valid_at: e.valid_at,
+                        invalid_at: e.invalid_at,
+                        attributes: properties(e)
+                    })[..$limit] AS matches
+                """
             )
-            + """ AS score
-            WHERE score > $min_score
-            WITH edge, e, score
-            ORDER BY score DESC
-            RETURN edge.uuid AS search_edge_uuid,
-                collect({
-                    uuid: e.uuid,
-                    source_node_uuid: startNode(e).uuid,
-                    target_node_uuid: endNode(e).uuid,
-                    created_at: e.created_at,
-                    name: e.name,
-                    group_id: e.group_id,
-                    fact: e.fact,
-                    fact_embedding: e.fact_embedding,
-                    episodes: e.episodes,
-                    expired_at: e.expired_at,
-                    valid_at: e.valid_at,
-                    invalid_at: e.invalid_at,
-                    attributes: properties(e)
-                })[..$limit] AS matches
-            """
-        )
 
         results, _, _ = await driver.execute_query(
             query,
@@ -1436,11 +1725,11 @@ async def get_edge_invalidation_candidates(
             limit=limit,
             min_score=min_score,
             routing_='r',
-            **query_params,
+            **filter_params,
         )
     invalidation_edges_dict: dict[str, list[EntityEdge]] = {
         result['search_edge_uuid']: [
-            get_entity_edge_from_record(record) for record in result['matches']
+            get_entity_edge_from_record(record, driver.provider) for record in result['matches']
         ]
         for result in results
     }
@@ -1479,13 +1768,21 @@ async def node_distance_reranker(
     filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
     scores: dict[str, float] = {center_node_uuid: 0.0}
 
-    # Find the shortest path to center node
-    results, header, _ = await driver.execute_query(
-        """
+    query = """
+    UNWIND $node_uuids AS node_uuid
+    MATCH (center:Entity {uuid: $center_uuid})-[:RELATES_TO]-(n:Entity {uuid: node_uuid})
+    RETURN 1 AS score, node_uuid AS uuid
+    """
+    if driver.provider == GraphProvider.KUZU:
+        query = """
         UNWIND $node_uuids AS node_uuid
-        MATCH (center:Entity {uuid: $center_uuid})-[:RELATES_TO]-(n:Entity {uuid: node_uuid})
+        MATCH (center:Entity {uuid: $center_uuid})-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(n:Entity {uuid: node_uuid})
         RETURN 1 AS score, node_uuid AS uuid
-        """,
+        """
+
+    # Find the shortest path to center node
+    results, header, _ = await driver.execute_query(
+        query,
         node_uuids=filtered_uuids,
         center_uuid=center_node_uuid,
         routing_='r',
@@ -1536,6 +1833,10 @@ async def episode_mentions_reranker(
     for result in results:
         scores[result['uuid']] = result['score']
 
+    for uuid in sorted_uuids:
+        if uuid not in scores:
+            scores[uuid] = float('inf')
+
     # rerank on shortest distance
     sorted_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
 
@@ -1667,13 +1968,23 @@ async def get_embeddings_for_edges(
             split(e.fact_embedding, ",") AS fact_embedding
         """
     else:
-        query = """
-        MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
+        match_query = """
+            MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity)-[:RELATES_TO]-(e:RelatesToNode_)-[:RELATES_TO]-(m:Entity)
+            """
+
+        query = (
+            match_query
+            + """
         WHERE e.uuid IN $edge_uuids
         RETURN DISTINCT
             e.uuid AS uuid,
             e.fact_embedding AS fact_embedding
         """
+        )
     results, _, _ = await driver.execute_query(
         query,
         edge_uuids=[edge.uuid for edge in edges],