graphiti-core 0.10.4__py3-none-any.whl → 0.11.0__py3-none-any.whl

graphiti_core/search/search_utils.py

@@ -26,7 +26,7 @@ from typing_extensions import LiteralString
 from graphiti_core.edges import EntityEdge, get_entity_edge_from_record
 from graphiti_core.helpers import (
     DEFAULT_DATABASE,
-    USE_PARALLEL_RUNTIME,
+    RUNTIME_QUERY,
     lucene_sanitize,
     normalize_l2,
     semaphore_gather,
@@ -54,20 +54,6 @@ DEFAULT_MMR_LAMBDA = 0.5
 MAX_SEARCH_DEPTH = 3
 MAX_QUERY_LENGTH = 32
 
-SEARCH_ENTITY_NODE_RETURN: LiteralString = """
-        OPTIONAL MATCH (e:Episodic)-[r:MENTIONS]->(n)
-        WITH n, score, collect(e.uuid) AS episodes
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.name_embedding AS name_embedding,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes,
-            episodes"""
-
 
 def fulltext_query(query: str, group_ids: list[str] | None = None):
     group_ids_filter_list = (
@@ -221,10 +207,6 @@ async def edge_similarity_search(
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityEdge]:
     # vector similarity search over embedded facts
-    runtime_query: LiteralString = (
-        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
-    )
-
     query_params: dict[str, Any] = {}
 
     filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
@@ -244,9 +226,10 @@ async def edge_similarity_search(
             group_filter_query += '\nAND (m.uuid IN [$source_uuid, $target_uuid])'
 
     query: LiteralString = (
-        """
-                                                                                                                                            MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
-                                                                                                                                            """
+        RUNTIME_QUERY
+        + """
+                                                                                                                                                                MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                                                                                                                                               """
         + group_filter_query
         + filter_query
         + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
@@ -270,7 +253,7 @@ async def edge_similarity_search(
     )
 
     records, _, _ = await driver.execute_query(
-        runtime_query + query,
+        query,
         query_params,
         search_vector=search_vector,
         source_uuid=source_node_uuid,
@@ -358,12 +341,12 @@ async def node_fulltext_search(
 
     query = (
         """
-                CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
-                YIELD node AS n, score
-                WHERE n:Entity
-                """
+                                CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
+                                YIELD node AS n, score
+                                WHERE n:Entity
+                                """
         + filter_query
-        + SEARCH_ENTITY_NODE_RETURN
+        + ENTITY_NODE_RETURN
         + """
         ORDER BY score DESC
         """
@@ -392,10 +375,6 @@ async def node_similarity_search(
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityNode]:
     # vector similarity search over entity names
-    runtime_query: LiteralString = (
-        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
-    )
-
     query_params: dict[str, Any] = {}
 
     group_filter_query: LiteralString = ''
@@ -407,7 +386,7 @@ async def node_similarity_search(
     query_params.update(filter_params)
 
     records, _, _ = await driver.execute_query(
-        runtime_query
+        RUNTIME_QUERY
         + """
             MATCH (n:Entity)
             """
@@ -416,7 +395,7 @@ async def node_similarity_search(
         + """
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
             WHERE score > $min_score"""
-        + SEARCH_ENTITY_NODE_RETURN
+        + ENTITY_NODE_RETURN
         + """
         ORDER BY score DESC
         LIMIT $limit
@@ -556,10 +535,6 @@ async def community_similarity_search(
     min_score=DEFAULT_MIN_SCORE,
 ) -> list[CommunityNode]:
     # vector similarity search over entity names
-    runtime_query: LiteralString = (
-        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
-    )
-
     query_params: dict[str, Any] = {}
 
     group_filter_query: LiteralString = ''
@@ -568,7 +543,7 @@ async def community_similarity_search(
         query_params['group_ids'] = group_ids
 
     records, _, _ = await driver.execute_query(
-        runtime_query
+        RUNTIME_QUERY
         + """
            MATCH (comm:Community)
            """
@@ -674,86 +649,219 @@ async def hybrid_node_search(
 
 async def get_relevant_nodes(
     driver: AsyncDriver,
-    search_filter: SearchFilters,
     nodes: list[EntityNode],
-) -> list[EntityNode]:
-    """
-    Retrieve relevant nodes based on the provided list of EntityNodes.
+    search_filter: SearchFilters,
+    min_score: float = DEFAULT_MIN_SCORE,
+    limit: int = RELEVANT_SCHEMA_LIMIT,
+) -> list[list[EntityNode]]:
+    if len(nodes) == 0:
+        return []
 
-    This method performs a hybrid search using both the names and embeddings
-    of the input nodes to find relevant nodes in the graph database.
+    group_id = nodes[0].group_id
 
-    Parameters
-    ----------
-    nodes : list[EntityNode]
-        A list of EntityNode objects to use as the basis for the search.
-    driver : AsyncDriver
-        The Neo4j driver instance for database operations.
+    # vector similarity search over entity names
+    query_params: dict[str, Any] = {}
 
-    Returns
-    -------
-    list[EntityNode]
-        A list of EntityNode objects that are deemed relevant based on the input nodes.
+    filter_query, filter_params = node_search_filter_query_constructor(search_filter)
+    query_params.update(filter_params)
 
-    Notes
-    -----
-    This method uses the hybrid_node_search function to perform the search,
-    which combines fulltext search and vector similarity search.
-    It extracts the names and name embeddings (if available) from the input nodes
-    to use as search criteria.
-    """
-    relevant_nodes = await hybrid_node_search(
-        [node.name for node in nodes],
-        [node.name_embedding for node in nodes if node.name_embedding is not None],
-        driver,
-        search_filter,
-        [node.group_id for node in nodes],
+    query = (
+        RUNTIME_QUERY
+        + """UNWIND $nodes AS node
+    MATCH (n:Entity {group_id: $group_id})
+            """
+        + filter_query
+        + """
+        WITH node, n, vector.similarity.cosine(n.name_embedding, node.name_embedding) AS score
+        WHERE score > $min_score
+        WITH node, collect(n)[..$limit] AS top_vector_nodes, collect(n.uuid) AS vector_node_uuids
+        
+        CALL db.index.fulltext.queryNodes("node_name_and_summary", 'group_id:"' + $group_id + '" AND ' + node.name, {limit: $limit}) 
+        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,
+        query_params,
+        nodes=[
+            {'uuid': node.uuid, 'name': node.name, 'name_embedding': node.name_embedding}
+            for node in nodes
+        ],
+        group_id=group_id,
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+
+    relevant_nodes_dict: dict[str, list[EntityNode]] = {
+        result['search_node_uuid']: [
+            get_entity_node_from_record(record) for record in result['matches']
+        ]
+        for result in results
+    }
+
+    relevant_nodes = [relevant_nodes_dict.get(node.uuid, []) for node in nodes]
+
     return relevant_nodes
 
 
 async def get_relevant_edges(
     driver: AsyncDriver,
     edges: list[EntityEdge],
-    source_node_uuid: str | None,
-    target_node_uuid: str | None,
+    search_filter: SearchFilters,
+    min_score: float = DEFAULT_MIN_SCORE,
     limit: int = RELEVANT_SCHEMA_LIMIT,
-) -> list[EntityEdge]:
-    start = time()
-    relevant_edges: list[EntityEdge] = []
-    relevant_edge_uuids = set()
-
-    results = await semaphore_gather(
-        *[
-            edge_similarity_search(
-                driver,
-                edge.fact_embedding,
-                source_node_uuid,
-                target_node_uuid,
-                SearchFilters(),
-                [edge.group_id],
-                limit,
-            )
-            for edge in edges
-            if edge.fact_embedding is not None
-        ]
-    )
+) -> list[list[EntityEdge]]:
+    if len(edges) == 0:
+        return []
 
-    for result in results:
-        for edge in result:
-            if edge.uuid in relevant_edge_uuids:
-                continue
+    query_params: dict[str, Any] = {}
 
-            relevant_edge_uuids.add(edge.uuid)
-            relevant_edges.append(edge)
+    filter_query, filter_params = edge_search_filter_query_constructor(search_filter)
+    query_params.update(filter_params)
 
-    end = time()
-    logger.debug(f'Found relevant edges: {relevant_edge_uuids} in {(end - start) * 1000} ms')
+    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, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) 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
+                })[..$limit] AS matches
+        """
+    )
+
+    results, _, _ = await driver.execute_query(
+        query,
+        query_params,
+        edges=[edge.model_dump() for edge in edges],
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+    relevant_edges_dict: dict[str, list[EntityEdge]] = {
+        result['search_edge_uuid']: [
+            get_entity_edge_from_record(record) for record in result['matches']
+        ]
+        for result in results
+    }
+
+    relevant_edges = [relevant_edges_dict.get(edge.uuid, []) for edge in edges]
 
     return relevant_edges
 
 
+async def get_edge_invalidation_candidates(
+    driver: AsyncDriver,
+    edges: list[EntityEdge],
+    search_filter: SearchFilters,
+    min_score: float = DEFAULT_MIN_SCORE,
+    limit: int = RELEVANT_SCHEMA_LIMIT,
+) -> list[list[EntityEdge]]:
+    if len(edges) == 0:
+        return []
+
+    query_params: dict[str, Any] = {}
+
+    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, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) 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
+                })[..$limit] AS matches
+        """
+    )
+
+    results, _, _ = await driver.execute_query(
+        query,
+        query_params,
+        edges=[edge.model_dump() for edge in edges],
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+    invalidation_edges_dict: dict[str, list[EntityEdge]] = {
+        result['search_edge_uuid']: [
+            get_entity_edge_from_record(record) for record in result['matches']
+        ]
+        for result in results
+    }
+
+    invalidation_edges = [invalidation_edges_dict.get(edge.uuid, []) for edge in edges]
+
+    return invalidation_edges
+
+
 # takes in a list of rankings of uuids
 def rrf(results: list[list[str]], rank_const=1, min_score: float = 0) -> list[str]:
     scores: dict[str, float] = defaultdict(float)

graphiti_core/utils/bulk_utils.py

@@ -26,6 +26,7 @@ from pydantic import BaseModel
 from typing_extensions import Any
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
+from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.models.edges.edge_db_queries import (
@@ -128,16 +129,18 @@ async def add_nodes_and_edges_bulk_tx(
 
     await tx.run(EPISODIC_NODE_SAVE_BULK, episodes=episodes)
     await tx.run(ENTITY_NODE_SAVE_BULK, nodes=nodes)
-    await tx.run(EPISODIC_EDGE_SAVE_BULK, episodic_edges=[dict(edge) for edge in episodic_edges])
-    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[dict(edge) for edge in entity_edges])
+    await tx.run(
+        EPISODIC_EDGE_SAVE_BULK, episodic_edges=[edge.model_dump() for edge in episodic_edges]
+    )
+    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[edge.model_dump() for edge in entity_edges])
 
 
 async def extract_nodes_and_edges_bulk(
-    llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
+    clients: GraphitiClients, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
     extracted_nodes_bulk = await semaphore_gather(
         *[
-            extract_nodes(llm_client, episode, previous_episodes)
+            extract_nodes(clients, episode, previous_episodes)
             for episode, previous_episodes in episode_tuples
         ]
     )
@@ -150,7 +153,7 @@ async def extract_nodes_and_edges_bulk(
     extracted_edges_bulk = await semaphore_gather(
         *[
             extract_edges(
-                llm_client,
+                clients,
                 episode,
                 extracted_nodes_bulk[i],
                 previous_episodes_list[i],
@@ -189,7 +192,7 @@ async def dedupe_nodes_bulk(
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
         await semaphore_gather(
-            *[get_relevant_nodes(driver, SearchFilters(), node_chunk) for node_chunk in node_chunks]
+            *[get_relevant_nodes(driver, node_chunk, SearchFilters()) for node_chunk in node_chunks]
         )
     )
 
@@ -223,7 +226,7 @@ async def dedupe_edges_bulk(
 
     relevant_edges_chunks: list[list[EntityEdge]] = list(
         await semaphore_gather(
-            *[get_relevant_edges(driver, edge_chunk, None, None) for edge_chunk in edge_chunks]
+            *[get_relevant_edges(driver, edge_chunk, SearchFilters()) for edge_chunk in edge_chunks]
         )
     )