graphiti-core 0.18.8__py3-none-any.whl → 0.19.0__py3-none-any.whl

graphiti_core/search/search_utils.py

@@ -15,6 +15,7 @@ limitations under the License.
 """
 
 import logging
+import os
 from collections import defaultdict
 from time import time
 from typing import Any
@@ -36,10 +37,13 @@ from graphiti_core.helpers import (
     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,
@@ -54,6 +58,7 @@ from graphiti_core.search.search_filters import (
 )
 
 logger = logging.getLogger(__name__)
+USE_HNSW = os.getenv('USE_HNSW', '').lower() in ('true', '1', 'yes')
 
 RELEVANT_SCHEMA_LIMIT = 10
 DEFAULT_MIN_SCORE = 0.6
@@ -62,9 +67,30 @@ MAX_SEARCH_DEPTH = 3
 MAX_QUERY_LENGTH = 128
 
 
-def fulltext_query(query: str, group_ids: list[str] | None = None, fulltext_syntax: str = ''):
+def calculate_cosine_similarity(vector1: list[float], vector2: list[float]) -> float:
+    """
+    Calculates the cosine similarity between two vectors using NumPy.
+    """
+    dot_product = np.dot(vector1, vector2)
+    norm_vector1 = np.linalg.norm(vector1)
+    norm_vector2 = np.linalg.norm(vector2)
+
+    if norm_vector1 == 0 or norm_vector2 == 0:
+        return 0  # Handle cases where one or both vectors are zero vectors
+
+    return dot_product / (norm_vector1 * norm_vector2)
+
+
+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:
@@ -108,12 +134,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
 
@@ -125,7 +151,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
         """
@@ -147,40 +173,105 @@ 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)
-
-    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 """
-        + filter_query
-        + """
-        WITH e, score, n, m
-        RETURN
-        """
-        + ENTITY_EDGE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
+    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)
         """
-    )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
     )
 
-    edges = [get_entity_edge_from_record(record) for record in records]
+    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
+        if res['hits']['total']['value'] > 0:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in res['hits']['hits']:
+                input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
+
+            # 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 
+                                    """
+                + 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,
+                    e.group_id AS group_id,
+                    n.uuid AS source_node_uuid,
+                    m.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
+                ORDER BY score DESC LIMIT $limit
+                            """
+            )
+
+            records, _, _ = await driver.execute_query(
+                query,
+                query=fuzzy_query,
+                ids=input_ids,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        query = (
+            get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
+            + match_query
+            + filter_query
+            + """
+            WITH e, score, n, m
+            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,
+        )
+
+    edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
     return edges
 
@@ -195,56 +286,130 @@ 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)
+            """
+            + filter_query
+            + """
+            RETURN DISTINCT id(e) as id, e.fact_embedding as embedding
+            """
+        )
+        resp, header, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
 
-    query = (
-        RUNTIME_QUERY
-        + """
-        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-        """
-        + group_filter_query
-        + filter_query
-        + """
-        WITH DISTINCT e, n, m, """
-        + get_vector_cosine_func_query('e.fact_embedding', '$search_vector', driver.provider)
-        + """ AS score
-        WHERE score > $min_score
-        RETURN
-        """
-        + ENTITY_EDGE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-    )
+        if len(resp) > 0:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in resp:
+                if r['embedding']:
+                    score = calculate_cosine_similarity(
+                        search_vector, list(map(float, r['embedding'].split(',')))
+                    )
+                    if score > min_score:
+                        input_ids.append({'id': r['id'], 'score': score})
+
+            # Match the edge ides and return the values
+            query = """
+                UNWIND $ids as i
+                MATCH ()-[r]->()
+                WHERE id(r) = i.id
+                RETURN
+                    r.uuid AS uuid,
+                    r.group_id AS group_id,
+                    startNode(r).uuid AS source_node_uuid,
+                    endNode(r).uuid AS target_node_uuid,
+                    r.created_at AS created_at,
+                    r.name AS name,
+                    r.fact AS fact,
+                    split(r.episodes, ",") AS episodes,
+                    r.expired_at AS expired_at,
+                    r.valid_at AS valid_at,
+                    r.invalid_at AS invalid_at,
+                    properties(r) AS attributes
+                ORDER BY i.score DESC
+                LIMIT $limit
+                    """
+            records, _, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                search_vector=search_vector,
+                limit=limit,
+                min_score=min_score,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        query = (
+            RUNTIME_QUERY
+            + match_query
+            + filter_query
+            + """
+            WITH DISTINCT e, n, m, """
+            + get_vector_cosine_func_query('e.fact_embedding', search_vector_var, driver.provider)
+            + """ AS score
+            WHERE score > $min_score
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider)
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        search_vector=search_vector,
-        limit=limit,
-        min_score=min_score,
-        routing_='r',
-        **query_params,
-    )
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            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
 
@@ -258,40 +423,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)
-
-    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)
-            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
-        """
+    filter_queries, filter_params = edge_search_filter_query_constructor(
+        search_filter, driver.provider
     )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        bfs_origin_node_uuids=bfs_origin_node_uuids,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
+    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: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:
+        if driver.provider == GraphProvider.NEPTUNE:
+            query = (
+                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 {{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,
+            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
 
@@ -302,39 +543,90 @@ 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)
 
-    query = (
-        get_nodes_query(driver.provider, 'node_name_and_summary', '$query')
-        + """
-        YIELD node AS n, score
-        WHERE n:Entity AND n.group_id IN $group_ids
-        """
-        + filter_query
-        + """
-        WITH n, score
-        ORDER BY score DESC
-        LIMIT $limit
-        RETURN
-        """
-        + ENTITY_NODE_RETURN
+    filter_queries, filter_params = node_search_filter_query_constructor(
+        search_filter, driver.provider
     )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
+    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
+        if res['hits']['total']['value'] > 0:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in res['hits']['hits']:
+                input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
+
+            # Match the edge ides and return the values
+            query = (
+                """
+                                    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
+                            """
+            )
+            records, _, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                query=fuzzy_query,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        index_name = (
+            'node_name_and_summary'
+            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(index_name, '$query', limit=limit, provider=driver.provider)
+            + yield_query
+            + filter_query
+            + """
+            WITH n, score
+            ORDER BY score DESC
+            LIMIT $limit
+            RETURN
+            """
+            + get_entity_node_return_query(driver.provider)
+        )
 
-    nodes = [get_entity_node_from_record(record) for record in records]
+        records, _, _ = await driver.execute_query(
+            query,
+            query=fuzzy_query,
+            limit=limit,
+            routing_='r',
+            **filter_params,
+        )
+
+    nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
     return nodes
 
@@ -346,49 +638,140 @@ 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 = ''
+    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)
+            """
+            + filter_query
+            + """
+            RETURN DISTINCT id(n) as id, n.name_embedding as embedding
+            """
+        )
+        resp, header, _ = await driver.execute_query(
+            query,
+            params=filter_params,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+        )
 
-    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
-    query_params.update(filter_params)
+        if len(resp) > 0:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in resp:
+                if r['embedding']:
+                    score = calculate_cosine_similarity(
+                        search_vector, list(map(float, r['embedding'].split(',')))
+                    )
+                    if score > min_score:
+                        input_ids.append({'id': r['id'], 'score': score})
+
+            # Match the edge ides and return the values
+            query = (
+                """
+                                        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
+                """
+            )
+            records, header, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                search_vector=search_vector,
+                limit=limit,
+                min_score=min_score,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    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 ''
+        )
+        query = (
+            f"""
+                    CALL db.index.vector.queryNodes('{index_name}', {limit}, $search_vector) YIELD node AS n, score
+                    """
+            + filter_query
+            + """
+                    AND score > $min_score
+                    RETURN
+                    """
+            + get_entity_node_return_query(driver.provider)
+            + """
+                    ORDER BY score DESC
+                    LIMIT $limit
+                    """
+        )
 
-    query = (
-        RUNTIME_QUERY
-        + """
-        MATCH (n:Entity)
-        """
-        + group_filter_query
-        + filter_query
-        + """
-        WITH n, """
-        + get_vector_cosine_func_query('n.name_embedding', '$search_vector', driver.provider)
-        + """ AS score
-        WHERE score > $min_score
-        RETURN
-        """
-        + ENTITY_NODE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-    )
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        search_vector=search_vector,
-        limit=limit,
-        min_score=min_score,
-        routing_='r',
-        **query_params,
-    )
+    else:
+        query = (
+            RUNTIME_QUERY
+            + """
+            MATCH (n:Entity)
+            """
+            + filter_query
+            + """
+            WITH n, """
+            + get_vector_cosine_func_query('n.name_embedding', search_vector_var, driver.provider)
+            + """ AS score
+            WHERE score > $min_score
+            RETURN
+            """
+            + get_entity_node_return_query(driver.provider)
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            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
 
@@ -401,38 +784,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))
 
-    query = (
+    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:
+        match_queries = [
+            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 {{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
+            """
+        ]
+
+    if driver.provider == GraphProvider.KUZU:
+        depth = bfs_max_depth * 2
+        match_queries = [
+            """
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH (origin:Episodic {uuid: origin_uuid})-[:MENTIONS]->(n:Entity)
             WHERE n.group_id = origin.group_id
-            AND origin.group_id IN $group_ids
-        """
-        + filter_query
-        + """
-        RETURN
-        """
-        + ENTITY_NODE_RETURN
-        + """
-        LIMIT $limit
-        """
-    )
+            """,
+            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:Episodic {{uuid: origin_uuid}})-[:MENTIONS]->(:Entity)-[:RELATES_TO*2..{depth}]->(n:Entity)
+                WHERE n.group_id = origin.group_id
+            """)
+
+    records = []
+    for match_query in match_queries:
+        sub_records, _, _ = await driver.execute_query(
+            match_query
+            + filter_query
+            + """
+            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
 
@@ -443,35 +870,84 @@ 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 []
 
-    query = (
-        get_nodes_query(driver.provider, 'episode_content', '$query')
-        + """
-        YIELD node AS episode, score
-        MATCH (e:Episodic)
-        WHERE e.uuid = episode.uuid
-        AND e.group_id IN $group_ids
-        RETURN
-        """
-        + EPISODIC_NODE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-    )
+    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:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in res['hits']['hits']:
+                input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
+
+            # Match the edge ides and return the values
+            query = """
+                UNWIND $ids as i
+                MATCH (e:Episodic)
+                WHERE e.uuid=i.id
+            RETURN
+                    e.content AS content,
+                    e.created_at AS created_at,
+                    e.valid_at AS valid_at,
+                    e.uuid AS uuid,
+                    e.name AS name,
+                    e.group_id AS group_id,
+                    e.source_description AS source_description,
+                    e.source AS source,
+                    e.entity_edges AS entity_edges
+                ORDER BY i.score DESC
+                LIMIT $limit
+            """
+            records, _, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                query=fuzzy_query,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        index_name = (
+            'episode_content'
+            if not use_local_indexes
+            else 'episode_content_'
+            + (group_ids[0].replace('-', '') if group_ids is not None else '')
+        )
+        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
+            """
+            + group_filter_query
+            + """
+            RETURN
+            """
+            + EPISODIC_NODE_RETURN
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query, query=fuzzy_query, limit=limit, routing_='r', **filter_params
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-    )
     episodes = [get_episodic_node_from_record(record) for record in records]
 
     return episodes
@@ -484,31 +960,75 @@ 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 []
 
-    query = (
-        get_nodes_query(driver.provider, 'community_name', '$query')
-        + """
-        YIELD node AS n, score
-        WHERE n.group_id IN $group_ids
-        RETURN
-        """
-        + COMMUNITY_NODE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-    )
+    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:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in res['hits']['hits']:
+                input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
+
+            # Match the edge ides and return the values
+            query = """
+                UNWIND $ids as i
+                MATCH (comm:Community)
+                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.summary AS summary,
+                    [x IN split(comm.name_embedding, ",") | toFloat(x)]AS name_embedding
+                ORDER BY i.score DESC
+                LIMIT $limit
+            """
+            records, _, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                query=fuzzy_query,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        query = (
+            get_nodes_query('community_name', '$query', limit=limit, provider=driver.provider)
+            + yield_query
+            + """
+            WITH c, score
+            """
+            + group_filter_query
+            + """
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query, query=fuzzy_query, limit=limit, routing_='r', **filter_params
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-    )
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -526,38 +1046,101 @@ 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
 
-    query = (
-        RUNTIME_QUERY
-        + """
-        MATCH (n:Community)
-        """
-        + group_filter_query
-        + """
-        WITH n,
-        """
-        + get_vector_cosine_func_query('n.name_embedding', '$search_vector', driver.provider)
-        + """ AS score
-        WHERE score > $min_score
-        RETURN
-        """
-        + COMMUNITY_NODE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
-        """
-    )
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = (
+            RUNTIME_QUERY
+            + """
+            MATCH (n:Community)
+            """
+            + group_filter_query
+            + """
+            RETURN DISTINCT id(n) as id, n.name_embedding as embedding
+            """
+        )
+        resp, header, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
+
+        if len(resp) > 0:
+            # Calculate Cosine similarity then return the edge ids
+            input_ids = []
+            for r in resp:
+                if r['embedding']:
+                    score = calculate_cosine_similarity(
+                        search_vector, list(map(float, r['embedding'].split(',')))
+                    )
+                    if score > min_score:
+                        input_ids.append({'id': r['id'], 'score': score})
+
+            # Match the edge ides and return the values
+            query = """
+                    UNWIND $ids as i
+                    MATCH (comm:Community)
+                    WHERE id(comm)=i.id
+                    RETURN
+                        comm.uuid As uuid,
+                        comm.group_id AS group_id,
+                        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
+                    LIMIT $limit
+                """
+            records, header, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                search_vector=search_vector,
+                limit=limit,
+                min_score=min_score,
+                routing_='r',
+                **query_params,
+            )
+        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 (c:Community)
+            """
+            + group_filter_query
+            + """
+            WITH c,
+            """
+            + get_vector_cosine_func_query('c.name_embedding', search_vector_var, driver.provider)
+            + """ AS score
+            WHERE score > $min_score
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        search_vector=search_vector,
-        limit=limit,
-        min_score=min_score,
-        routing_='r',
-        **query_params,
-    )
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -648,67 +1231,129 @@ 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 = (
+            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',
+                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 = (
+            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(
+                '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,
@@ -716,12 +1361,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
     }
@@ -741,25 +1386,53 @@ 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})
+            """
+            + filter_query
+            + """
+            WITH e, edge
+            RETURN DISTINCT id(e) as id, e.fact_embedding as source_embedding, edge.uuid as search_edge_uuid,
+            edge.fact_embedding as target_embedding
+            """
+        )
+        resp, _, _ = await driver.execute_query(
+            query,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
 
-    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)
-        + """ AS score
-        WHERE score > $min_score
-        WITH edge, e, score
-        ORDER BY score DESC
+        # Calculate Cosine similarity then return the edge ids
+        input_ids = []
+        for r in resp:
+            score = calculate_cosine_similarity(
+                list(map(float, r['source_embedding'].split(','))), r['target_embedding']
+            )
+            if score > min_score:
+                input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']})
+
+        # Match the edge ides and return the values
+        query = """
+        UNWIND $ids AS edge
+        MATCH ()-[e]->()
+        WHERE id(e) = edge.id
+        WITH edge, e
+        ORDER BY edge.score DESC
         RETURN edge.uuid AS search_edge_uuid,
             collect({
                 uuid: e.uuid,
@@ -769,28 +1442,119 @@ async def get_relevant_edges(
                 name: e.name,
                 group_id: e.group_id,
                 fact: e.fact,
-                fact_embedding: e.fact_embedding,
-                episodes: e.episodes,
+                fact_embedding: [x IN split(e.fact_embedding, ",") | toFloat(x)],
+                episodes: split(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,
-        edges=[edge.model_dump() for edge in edges],
-        limit=limit,
-        min_score=min_score,
-        routing_='r',
-        **query_params,
-    )
+                """
+
+        results, _, _ = await driver.execute_query(
+            query,
+            ids=input_ids,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
+    else:
+        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 = (
+                RUNTIME_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 = (
+                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
+                )
+                + """ 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,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **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
     }
@@ -810,26 +1574,55 @@ 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]
+            """
+            + filter_query
+            + """
+            WITH e, edge
+            RETURN DISTINCT id(e) as id, e.fact_embedding as source_embedding,
+            edge.fact_embedding as target_embedding,
+            edge.uuid as search_edge_uuid
+            """
+        )
+        resp, _, _ = await driver.execute_query(
+            query,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
 
-    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)
-        + """ AS score
-        WHERE score > $min_score
-        WITH edge, e, score
-        ORDER BY score DESC
+        # Calculate Cosine similarity then return the edge ids
+        input_ids = []
+        for r in resp:
+            score = calculate_cosine_similarity(
+                list(map(float, r['source_embedding'].split(','))), r['target_embedding']
+            )
+            if score > min_score:
+                input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']})
+
+        # Match the edge ides and return the values
+        query = """
+        UNWIND $ids AS edge
+        MATCH ()-[e]->()
+        WHERE id(e) = edge.id
+        WITH edge, e
+        ORDER BY edge.score DESC
         RETURN edge.uuid AS search_edge_uuid,
             collect({
                 uuid: e.uuid,
@@ -839,27 +1632,119 @@ async def get_edge_invalidation_candidates(
                 name: e.name,
                 group_id: e.group_id,
                 fact: e.fact,
-                fact_embedding: e.fact_embedding,
-                episodes: e.episodes,
+                fact_embedding: [x IN split(e.fact_embedding, ",") | toFloat(x)],
+                episodes: split(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,
-        edges=[edge.model_dump() for edge in edges],
-        limit=limit,
-        min_score=min_score,
-        routing_='r',
-        **query_params,
-    )
+                """
+        results, _, _ = await driver.execute_query(
+            query,
+            ids=input_ids,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **filter_params,
+        )
+    else:
+        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 = (
+                RUNTIME_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 = (
+                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
+                )
+                + """ 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,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **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
     }
@@ -898,13 +1783,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',
@@ -955,6 +1848,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])
 
@@ -1007,14 +1904,24 @@ def maximal_marginal_relevance(
 async def get_embeddings_for_nodes(
     driver: GraphDriver, nodes: list[EntityNode]
 ) -> dict[str, list[float]]:
-    results, _, _ = await driver.execute_query(
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = """
+        MATCH (n:Entity)
+        WHERE n.uuid IN $node_uuids
+        RETURN DISTINCT
+            n.uuid AS uuid,
+            split(n.name_embedding, ",") AS name_embedding
         """
+    else:
+        query = """
         MATCH (n:Entity)
         WHERE n.uuid IN $node_uuids
         RETURN DISTINCT
             n.uuid AS uuid,
             n.name_embedding AS name_embedding
-        """,
+        """
+    results, _, _ = await driver.execute_query(
+        query,
         node_uuids=[node.uuid for node in nodes],
         routing_='r',
     )
@@ -1032,14 +1939,24 @@ async def get_embeddings_for_nodes(
 async def get_embeddings_for_communities(
     driver: GraphDriver, communities: list[CommunityNode]
 ) -> dict[str, list[float]]:
-    results, _, _ = await driver.execute_query(
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = """
+        MATCH (c:Community)
+        WHERE c.uuid IN $community_uuids
+        RETURN DISTINCT
+            c.uuid AS uuid,
+            split(c.name_embedding, ",") AS name_embedding
         """
+    else:
+        query = """
         MATCH (c:Community)
         WHERE c.uuid IN $community_uuids
         RETURN DISTINCT
             c.uuid AS uuid,
             c.name_embedding AS name_embedding
-        """,
+        """
+    results, _, _ = await driver.execute_query(
+        query,
         community_uuids=[community.uuid for community in communities],
         routing_='r',
     )
@@ -1057,14 +1974,34 @@ async def get_embeddings_for_communities(
 async def get_embeddings_for_edges(
     driver: GraphDriver, edges: list[EntityEdge]
 ) -> dict[str, list[float]]:
-    results, _, _ = await driver.execute_query(
-        """
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = """
         MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
         WHERE e.uuid IN $edge_uuids
+        RETURN DISTINCT
+            e.uuid AS uuid,
+            split(e.fact_embedding, ",") AS fact_embedding
+        """
+    else:
+        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],
         routing_='r',
     )