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

graphiti_core/nodes.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
@@ -32,10 +33,10 @@ from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_RETURN,
     COMMUNITY_NODE_RETURN_NEPTUNE,
-    ENTITY_NODE_RETURN,
     EPISODIC_NODE_RETURN,
     EPISODIC_NODE_RETURN_NEPTUNE,
     get_community_node_save_query,
+    get_entity_node_return_query,
     get_entity_node_save_query,
     get_episode_node_save_query,
 )
@@ -95,12 +96,37 @@ class Node(BaseModel, ABC):
             case GraphProvider.NEO4J:
                 await driver.execute_query(
                     """
-                MATCH (n:Entity|Episodic|Community {uuid: $uuid})
-                DETACH DELETE n
-                """,
+                    MATCH (n:Entity|Episodic|Community {uuid: $uuid})
+                    DETACH DELETE n
+                    """,
+                    uuid=self.uuid,
+                )
+            case GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{uuid: $uuid}})
+                        DETACH DELETE n
+                        """,
+                        uuid=self.uuid,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {uuid: $uuid})-[:RELATES_TO]->(e:RelatesToNode_)
+                    DETACH DELETE e
+                    """,
                     uuid=self.uuid,
                 )
-            case _:  # FalkorDB and Neptune
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {uuid: $uuid})
+                    DETACH DELETE n
+                    """,
+                    uuid=self.uuid,
+                )
+            case _:  # FalkorDB, Neptune
                 for label in ['Entity', 'Episodic', 'Community']:
                     await driver.execute_query(
                         f"""
@@ -136,8 +162,32 @@ class Node(BaseModel, ABC):
                         group_id=group_id,
                         batch_size=batch_size,
                     )
-
-            case _:  # FalkorDB and Neptune
+            case GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{group_id: $group_id}})
+                        DETACH DELETE n
+                        """,
+                        group_id=group_id,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {group_id: $group_id})-[:RELATES_TO]->(e:RelatesToNode_)
+                    DETACH DELETE e
+                    """,
+                    group_id=group_id,
+                )
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {group_id: $group_id})
+                    DETACH DELETE n
+                    """,
+                    group_id=group_id,
+                )
+            case _:  # FalkorDB, Neptune
                 for label in ['Entity', 'Episodic', 'Community']:
                     await driver.execute_query(
                         f"""
@@ -149,30 +199,59 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str], batch_size: int = 100):
-        if driver.provider == GraphProvider.FALKORDB:
-            for label in ['Entity', 'Episodic', 'Community']:
+        match driver.provider:
+            case GraphProvider.FALKORDB:
+                for label in ['Entity', 'Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label})
+                        WHERE n.uuid IN $uuids
+                        DETACH DELETE n
+                        """,
+                        uuids=uuids,
+                    )
+            case GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label})
+                        WHERE n.uuid IN $uuids
+                        DETACH DELETE n
+                        """,
+                        uuids=uuids,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
                 await driver.execute_query(
-                    f"""
-                       MATCH (n:{label})
-                       WHERE n.uuid IN $uuids
-                       DETACH DELETE n
-                       """,
+                    """
+                    MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_)
+                    WHERE n.uuid IN $uuids
+                    DETACH DELETE e
+                    """,
                     uuids=uuids,
                 )
-        else:
-            async with driver.session() as session:
-                await session.run(
+                await driver.execute_query(
                     """
-                    MATCH (n:Entity|Episodic|Community)
+                    MATCH (n:Entity)
                     WHERE n.uuid IN $uuids
-                    CALL {
-                        WITH n
-                        DETACH DELETE n
-                    } IN TRANSACTIONS OF $batch_size ROWS
+                    DETACH DELETE n
                     """,
                     uuids=uuids,
-                    batch_size=batch_size,
                 )
+            case _:  # Neo4J, Neptune
+                async with driver.session() as session:
+                    await session.run(
+                        """
+                        MATCH (n:Entity|Episodic|Community)
+                        WHERE n.uuid IN $uuids
+                        CALL {
+                            WITH n
+                            DETACH DELETE n
+                        } IN TRANSACTIONS OF $batch_size ROWS
+                        """,
+                        uuids=uuids,
+                        batch_size=batch_size,
+                    )
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
@@ -207,18 +286,24 @@ class EpisodicNode(Node):
                     }
                 ],
             )
+
+        episode_args = {
+            'uuid': self.uuid,
+            'name': self.name,
+            'group_id': self.group_id,
+            'source_description': self.source_description,
+            'content': self.content,
+            'entity_edges': self.entity_edges,
+            'created_at': self.created_at,
+            'valid_at': self.valid_at,
+            'source': self.source.value,
+        }
+
+        if driver.provider == GraphProvider.NEO4J:
+            episode_args['group_label'] = 'Episodic_' + self.group_id.replace('-', '')
+
         result = await driver.execute_query(
-            get_episode_node_save_query(driver.provider),
-            uuid=self.uuid,
-            name=self.name,
-            group_id=self.group_id,
-            group_label='Episodic_' + self.group_id.replace('-', ''),
-            source_description=self.source_description,
-            content=self.content,
-            entity_edges=self.entity_edges,
-            created_at=self.created_at,
-            valid_at=self.valid_at,
-            source=self.source.value,
+            get_episode_node_save_query(driver.provider), **episode_args
         )
 
         logger.debug(f'Saved Node to Graph: {self.uuid}')
@@ -376,17 +461,25 @@ class EntityNode(Node):
             'summary': self.summary,
             'created_at': self.created_at,
         }
-        entity_data.update(self.attributes or {})
 
-        if driver.provider == GraphProvider.NEPTUNE:
-            driver.save_to_aoss('node_name_and_summary', [entity_data])  # pyright: ignore reportAttributeAccessIssue
+        if driver.provider == GraphProvider.KUZU:
+            entity_data['attributes'] = json.dumps(self.attributes)
+            entity_data['labels'] = list(set(self.labels + ['Entity']))
+            result = await driver.execute_query(
+                get_entity_node_save_query(driver.provider, labels=''),
+                **entity_data,
+            )
+        else:
+            entity_data.update(self.attributes or {})
+            labels = ':'.join(self.labels + ['Entity', 'Entity_' + self.group_id.replace('-', '')])
 
-        labels = ':'.join(self.labels + ['Entity', 'Entity_' + self.group_id.replace('-', '')])
+            if driver.provider == GraphProvider.NEPTUNE:
+                driver.save_to_aoss('node_name_and_summary', [entity_data])  # pyright: ignore reportAttributeAccessIssue
 
-        result = await driver.execute_query(
-            get_entity_node_save_query(driver.provider, labels),
-            entity_data=entity_data,
-        )
+            result = await driver.execute_query(
+                get_entity_node_save_query(driver.provider, labels),
+                entity_data=entity_data,
+            )
 
         logger.debug(f'Saved Node to Graph: {self.uuid}')
 
@@ -399,12 +492,12 @@ class EntityNode(Node):
             MATCH (n:Entity {uuid: $uuid})
             RETURN
             """
-            + ENTITY_NODE_RETURN,
+            + get_entity_node_return_query(driver.provider),
             uuid=uuid,
             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]
 
         if len(nodes) == 0:
             raise NodeNotFoundError(uuid)
@@ -419,12 +512,12 @@ class EntityNode(Node):
             WHERE n.uuid IN $uuids
             RETURN
             """
-            + ENTITY_NODE_RETURN,
+            + get_entity_node_return_query(driver.provider),
             uuids=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
 
@@ -456,7 +549,7 @@ class EntityNode(Node):
             + """
             RETURN
             """
-            + ENTITY_NODE_RETURN
+            + get_entity_node_return_query(driver.provider)
             + with_embeddings_query
             + """
             ORDER BY n.uuid DESC
@@ -468,7 +561,7 @@ class EntityNode(Node):
             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
 
@@ -533,7 +626,7 @@ class CommunityNode(Node):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community {uuid: $uuid})
+            MATCH (c:Community {uuid: $uuid})
             RETURN
             """
             + (
@@ -556,8 +649,8 @@ class CommunityNode(Node):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community)
-            WHERE n.uuid IN $uuids
+            MATCH (c:Community)
+            WHERE c.uuid IN $uuids
             RETURN
             """
             + (
@@ -581,13 +674,13 @@ class CommunityNode(Node):
         limit: int | None = None,
         uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND n.uuid < $uuid' if uuid_cursor else ''
+        cursor_query: LiteralString = 'AND c.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community)
-            WHERE n.group_id IN $group_ids
+            MATCH (c:Community)
+            WHERE c.group_id IN $group_ids
             """
             + cursor_query
             + """
@@ -599,7 +692,7 @@ class CommunityNode(Node):
                 else COMMUNITY_NODE_RETURN
             )
             + """
-            ORDER BY n.uuid DESC
+            ORDER BY c.uuid DESC
             """
             + limit_query,
             group_ids=group_ids,
@@ -636,25 +729,35 @@ def get_episodic_node_from_record(record: Any) -> EpisodicNode:
     )
 
 
-def get_entity_node_from_record(record: Any) -> EntityNode:
+def get_entity_node_from_record(record: Any, provider: GraphProvider) -> EntityNode:
+    if provider == GraphProvider.KUZU:
+        attributes = json.loads(record['attributes']) if record['attributes'] else {}
+    else:
+        attributes = record['attributes']
+        attributes.pop('uuid', None)
+        attributes.pop('name', None)
+        attributes.pop('group_id', None)
+        attributes.pop('name_embedding', None)
+        attributes.pop('summary', None)
+        attributes.pop('created_at', None)
+        attributes.pop('labels', None)
+
+    labels = record.get('labels', [])
+    group_id = record.get('group_id')
+    if 'Entity_' + group_id.replace('-', '') in labels:
+        labels.remove('Entity_' + group_id.replace('-', ''))
+
     entity_node = EntityNode(
         uuid=record['uuid'],
         name=record['name'],
         name_embedding=record.get('name_embedding'),
-        group_id=record['group_id'],
-        labels=record['labels'],
+        group_id=group_id,
+        labels=labels,
         created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
-        attributes=record['attributes'],
+        attributes=attributes,
     )
 
-    entity_node.attributes.pop('uuid', None)
-    entity_node.attributes.pop('name', None)
-    entity_node.attributes.pop('group_id', None)
-    entity_node.attributes.pop('name_embedding', None)
-    entity_node.attributes.pop('summary', None)
-    entity_node.attributes.pop('created_at', None)
-
     return entity_node
 
 

graphiti_core/search/search.py

@@ -325,12 +325,20 @@ async def node_search(
     search_tasks = []
     if NodeSearchMethod.bm25 in config.search_methods:
         search_tasks.append(
-            node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
+            node_fulltext_search(
+                driver, query, search_filter, group_ids, 2 * limit, config.use_local_indexes
+            )
         )
     if NodeSearchMethod.cosine_similarity in config.search_methods:
         search_tasks.append(
             node_similarity_search(
-                driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
+                driver,
+                query_vector,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+                config.use_local_indexes,
             )
         )
     if NodeSearchMethod.bfs in config.search_methods:
@@ -426,7 +434,9 @@ async def episode_search(
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
-                episode_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
+                episode_fulltext_search(
+                    driver, query, search_filter, group_ids, 2 * limit, config.use_local_indexes
+                ),
             ]
         )
     )

graphiti_core/search/search_config.py

@@ -24,6 +24,7 @@ from graphiti_core.search.search_utils import (
     DEFAULT_MIN_SCORE,
     DEFAULT_MMR_LAMBDA,
     MAX_SEARCH_DEPTH,
+    USE_HNSW,
 )
 
 DEFAULT_SEARCH_LIMIT = 10
@@ -91,6 +92,7 @@ class NodeSearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class EpisodeSearchConfig(BaseModel):
@@ -99,6 +101,7 @@ class EpisodeSearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class CommunitySearchConfig(BaseModel):
@@ -107,6 +110,7 @@ class CommunitySearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class SearchConfig(BaseModel):

graphiti_core/search/search_filters.py

@@ -20,6 +20,8 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
+from graphiti_core.driver.driver import GraphProvider
+
 
 class ComparisonOperator(Enum):
     equals = '='
@@ -54,16 +56,21 @@ class SearchFilters(BaseModel):
 
 def node_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[str, dict[str, Any]]:
-    filter_query: str = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
     if filters.node_labels is not None:
-        node_labels = '|'.join(filters.node_labels)
-        node_label_filter = ' AND n:' + node_labels
-        filter_query += node_label_filter
+        if provider == GraphProvider.KUZU:
+            node_label_filter = 'list_has_all(n.labels, $labels)'
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels
+        filter_queries.append(node_label_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params
 
 
 def date_filter_query_constructor(
@@ -81,23 +88,29 @@ def date_filter_query_constructor(
 
 def edge_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[str, dict[str, Any]]:
-    filter_query: str = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
     if filters.edge_types is not None:
         edge_types = filters.edge_types
-        edge_types_filter = '\nAND e.name in $edge_types'
-        filter_query += edge_types_filter
+        filter_queries.append('e.name in $edge_types')
         filter_params['edge_types'] = edge_types
 
     if filters.node_labels is not None:
-        node_labels = '|'.join(filters.node_labels)
-        node_label_filter = '\nAND n:' + node_labels + ' AND m:' + node_labels
-        filter_query += node_label_filter
+        if provider == GraphProvider.KUZU:
+            node_label_filter = (
+                'list_has_all(n.labels, $labels) AND list_has_all(m.labels, $labels)'
+            )
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels + ' AND m:' + node_labels
+        filter_queries.append(node_label_filter)
 
     if filters.valid_at is not None:
-        valid_at_filter = '\nAND ('
+        valid_at_filter = '('
         for i, or_list in enumerate(filters.valid_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -125,10 +138,10 @@ def edge_search_filter_query_constructor(
             else:
                 valid_at_filter += ' OR '
 
-        filter_query += valid_at_filter
+        filter_queries.append(valid_at_filter)
 
     if filters.invalid_at is not None:
-        invalid_at_filter = ' AND ('
+        invalid_at_filter = '('
         for i, or_list in enumerate(filters.invalid_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -156,10 +169,10 @@ def edge_search_filter_query_constructor(
             else:
                 invalid_at_filter += ' OR '
 
-        filter_query += invalid_at_filter
+        filter_queries.append(invalid_at_filter)
 
     if filters.created_at is not None:
-        created_at_filter = ' AND ('
+        created_at_filter = '('
         for i, or_list in enumerate(filters.created_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -187,10 +200,10 @@ def edge_search_filter_query_constructor(
             else:
                 created_at_filter += ' OR '
 
-        filter_query += created_at_filter
+        filter_queries.append(created_at_filter)
 
     if filters.expired_at is not None:
-        expired_at_filter = ' AND ('
+        expired_at_filter = '('
         for i, or_list in enumerate(filters.expired_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -218,6 +231,6 @@ def edge_search_filter_query_constructor(
             else:
                 expired_at_filter += ' OR '
 
-        filter_query += expired_at_filter
+        filter_queries.append(expired_at_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params