graphiti-core 0.18.8__py3-none-any.whl → 0.19.0rc1__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
@@ -54,6 +55,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,6 +64,20 @@ MAX_SEARCH_DEPTH = 3
 MAX_QUERY_LENGTH = 128
 
 
+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 = None, fulltext_syntax: str = ''):
     group_ids_filter_list = (
         [fulltext_syntax + f'group_id:"{g}"' for g in group_ids] if group_ids is not None else []
@@ -153,32 +169,80 @@ async def edge_fulltext_search(
 
     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
-        """
-    )
+    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
+                + """
+                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,
+                group_ids=group_ids,
+                ids=input_ids,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            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 """
+            + filter_query
+            + """
+            WITH e, score, n, m
+            RETURN
+            """
+            + ENTITY_EDGE_RETURN
+            + """
+            ORDER BY score DESC
+            LIMIT $limit
+            """
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
+        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]
 
@@ -214,35 +278,100 @@ async def edge_similarity_search(
             query_params['target_uuid'] = target_node_uuid
             group_filter_query += '\nAND (m.uuid = $target_uuid)'
 
-    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 driver.provider == GraphProvider.NEPTUNE:
+        query = (
+            RUNTIME_QUERY
+            + """
+            MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)            
+            """
+            + group_filter_query
+            + 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',
+            **query_params,
+        )
 
-    records, _, _ = 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 ()-[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',
+                **query_params,
+            )
+        else:
+            return []
+    else:
+        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
+            """
+        )
+
+        records, _, _ = await driver.execute_query(
+            query,
+            search_vector=search_vector,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
 
     edges = [get_entity_edge_from_record(record) for record in records]
 
@@ -263,28 +392,58 @@ async def edge_bfs_search(
 
     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
-        """
-    )
+    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]-(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
+                """
+        )
+    else:
+        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
+            """
+        )
 
     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',
@@ -309,30 +468,70 @@ async def node_fulltext_search(
         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
-    )
+    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
+                                """
+                + ENTITY_NODE_RETURN
+                + """
+                ORDER BY i.score DESC
+                LIMIT $limit
+                            """
+            )
+            records, _, _ = await driver.execute_query(
+                query,
+                ids=input_ids,
+                query=fuzzy_query,
+                group_ids=group_ids,
+                limit=limit,
+                routing_='r',
+                **filter_params,
+            )
+        else:
+            return []
+    else:
+        index_name = (
+            'node_name_and_summary'
+            if not USE_HNSW
+            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
+          """
+            + filter_query
+            + """
+          WITH n, score
+          ORDER BY score DESC
+          LIMIT $limit
+          RETURN
+          """
+            + ENTITY_NODE_RETURN
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-        **filter_params,
-    )
+        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]
 
@@ -358,35 +557,124 @@ async def node_similarity_search(
     filter_query, filter_params = node_search_filter_query_constructor(search_filter)
     query_params.update(filter_params)
 
-    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
-        """
-    )
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = (
+            RUNTIME_QUERY
+            + """
+            MATCH (n:Entity)
+            """
+            + group_filter_query
+            + filter_query
+            + """
+            RETURN DISTINCT id(n) as id, n.name_embedding as embedding
+            """
+        )
+        resp, header, _ = await driver.execute_query(
+            query,
+            params=query_params,
+            search_vector=search_vector,
+            group_ids=group_ids,
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+        )
 
-    records, _, _ = 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 (n:Entity)
+                                    WHERE id(n)=i.id
+                                    RETURN 
+                                    """
+                + ENTITY_NODE_RETURN
+                + """
+                    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 []
+    elif driver.provider == GraphProvider.NEO4J and USE_HNSW:
+        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
+                    """
+            + group_filter_query
+            + filter_query
+            + """
+                    AND 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',
+            **query_params,
+        )
+
+    else:
+        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',
+            **query_params,
+        )
 
     nodes = [get_entity_node_from_record(record) for record in records]
 
@@ -407,22 +695,40 @@ async def node_bfs_search(
 
     filter_query, filter_params = node_search_filter_query_constructor(search_filter)
 
-    query = (
-        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)
-            WHERE n.group_id = origin.group_id
-            AND origin.group_id IN $group_ids
-        """
-        + filter_query
-        + """
-        RETURN
-        """
-        + ENTITY_NODE_RETURN
-        + """
-        LIMIT $limit
-        """
-    )
+    if driver.provider == GraphProvider.NEPTUNE:
+        query = (
+            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
+            """
+            + ENTITY_NODE_RETURN
+            + """
+            LIMIT $limit
+            """
+        )
+    else:
+        query = (
+            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)
+                WHERE n.group_id = origin.group_id
+                AND origin.group_id IN $group_ids
+            """
+            + filter_query
+            + """
+            RETURN
+            """
+            + ENTITY_NODE_RETURN
+            + """
+            LIMIT $limit
+            """
+        )
 
     records, _, _ = await driver.execute_query(
         query,
@@ -449,29 +755,72 @@ async def episode_fulltext_search(
     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
-        """
-    )
+    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,
+                group_ids=group_ids,
+                limit=limit,
+                routing_='r',
+            )
+        else:
+            return []
+    else:
+        index_name = (
+            'episode_content'
+            if not USE_HNSW
+            else 'episode_content_'
+            + (group_ids[0].replace('-', '') if group_ids is not None else '')
+        )
+        query = (
+            get_nodes_query(driver.provider, index_name, '$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
+            """
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-    )
+        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
@@ -488,27 +837,61 @@ async def community_fulltext_search(
     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
-        """
-    )
+    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,
+                group_ids=group_ids,
+                limit=limit,
+                routing_='r',
+            )
+        else:
+            return []
+    else:
+        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
+            """
+        )
 
-    records, _, _ = await driver.execute_query(
-        query,
-        query=fuzzy_query,
-        group_ids=group_ids,
-        limit=limit,
-        routing_='r',
-    )
+        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
@@ -529,35 +912,93 @@ async def community_similarity_search(
         group_filter_query += 'WHERE n.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,
+        )
 
-    records, _, _ = 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:
+        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
+            """
+        )
+
+        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
@@ -746,20 +1187,45 @@ async def get_relevant_edges(
     filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
     query_params.update(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
+    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',
+            **query_params,
+        )
+
+        # 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,24 +1235,69 @@ 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,
+            params=query_params,
+            ids=input_ids,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_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
+            )
+            + """ 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',
-        **query_params,
-    )
+        results, _, _ = await driver.execute_query(
+            query,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
 
     relevant_edges_dict: dict[str, list[EntityEdge]] = {
         result['search_edge_uuid']: [
@@ -815,21 +1326,47 @@ async def get_edge_invalidation_candidates(
     filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
     query_params.update(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
+    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',
+            **query_params,
+        )
+
+        # 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,24 +1376,68 @@ 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,
+            ids=input_ids,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_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
+            )
+            + """ 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',
-        **query_params,
-    )
+        results, _, _ = await driver.execute_query(
+            query,
+            edges=[edge.model_dump() for edge in edges],
+            limit=limit,
+            min_score=min_score,
+            routing_='r',
+            **query_params,
+        )
     invalidation_edges_dict: dict[str, list[EntityEdge]] = {
         result['search_edge_uuid']: [
             get_entity_edge_from_record(record) for record in result['matches']
@@ -1007,14 +1588,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 +1623,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 +1658,24 @@ 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:
+        query = """
         MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
         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',
     )