graphiti-core 0.17.4__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/utils/maintenance/graph_data_operations.py

@@ -15,53 +15,22 @@ limitations under the License.
 """
 
 import logging
-from datetime import datetime, timezone
+from datetime import datetime
 
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver
-from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices
-from graphiti_core.helpers import parse_db_date, semaphore_gather
-from graphiti_core.nodes import EpisodeType, EpisodicNode
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
+from graphiti_core.models.nodes.node_db_queries import (
+    EPISODIC_NODE_RETURN,
+    EPISODIC_NODE_RETURN_NEPTUNE,
+)
+from graphiti_core.nodes import EpisodeType, EpisodicNode, get_episodic_node_from_record
 
 EPISODE_WINDOW_LEN = 3
 
 logger = logging.getLogger(__name__)
 
 
-async def build_indices_and_constraints(driver: GraphDriver, delete_existing: bool = False):
-    if delete_existing:
-        records, _, _ = await driver.execute_query(
-            """
-        SHOW INDEXES YIELD name
-        """,
-        )
-        index_names = [record['name'] for record in records]
-        await semaphore_gather(
-            *[
-                driver.execute_query(
-                    """DROP INDEX $name""",
-                    name=name,
-                )
-                for name in index_names
-            ]
-        )
-    range_indices: list[LiteralString] = get_range_indices(driver.provider)
-
-    fulltext_indices: list[LiteralString] = get_fulltext_indices(driver.provider)
-
-    index_queries: list[LiteralString] = range_indices + fulltext_indices
-
-    await semaphore_gather(
-        *[
-            driver.execute_query(
-                query,
-            )
-            for query in index_queries
-        ]
-    )
-
-
 async def clear_data(driver: GraphDriver, group_ids: list[str] | None = None):
     async with driver.session() as session:
 
@@ -69,10 +38,19 @@ async def clear_data(driver: GraphDriver, group_ids: list[str] | None = None):
             await tx.run('MATCH (n) DETACH DELETE n')
 
         async def delete_group_ids(tx):
-            await tx.run(
-                'MATCH (n:Entity|Episodic|Community) WHERE n.group_id IN $group_ids DETACH DELETE n',
-                group_ids=group_ids,
-            )
+            labels = ['Entity', 'Episodic', 'Community']
+            if driver.provider == GraphProvider.KUZU:
+                labels.append('RelatesToNode_')
+
+            for label in labels:
+                await tx.run(
+                    f"""
+                    MATCH (n:{label})
+                    WHERE n.group_id IN $group_ids
+                    DETACH DELETE n
+                    """,
+                    group_ids=group_ids,
+                )
 
         if group_ids is None:
             await session.execute_write(delete_all)
@@ -101,26 +79,32 @@ async def retrieve_episodes(
     Returns:
         list[EpisodicNode]: A list of EpisodicNode objects representing the retrieved episodes.
     """
-    group_id_filter: LiteralString = (
-        '\nAND e.group_id IN $group_ids' if group_ids and len(group_ids) > 0 else ''
-    )
-    source_filter: LiteralString = '\nAND e.source = $source' if source is not None else ''
+
+    query_params: dict = {}
+    query_filter = ''
+    if group_ids and len(group_ids) > 0:
+        query_filter += '\nAND e.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
+    if source is not None:
+        query_filter += '\nAND e.source = $source'
+        query_params['source'] = source.name
 
     query: LiteralString = (
         """
-                                MATCH (e:Episodic) WHERE e.valid_at <= $reference_time
-                                """
-        + group_id_filter
-        + source_filter
+                                    MATCH (e:Episodic)
+                                    WHERE e.valid_at <= $reference_time
+                                    """
+        + query_filter
+        + """
+        RETURN
+        """
+        + (
+            EPISODIC_NODE_RETURN_NEPTUNE
+            if driver.provider == GraphProvider.NEPTUNE
+            else EPISODIC_NODE_RETURN
+        )
         + """
-        RETURN e.content AS content,
-            e.created_at AS created_at,
-            e.valid_at AS valid_at,
-            e.uuid AS uuid,
-            e.group_id AS group_id,
-            e.name AS name,
-            e.source_description AS source_description,
-            e.source AS source
         ORDER BY e.valid_at DESC
         LIMIT $num_episodes
         """
@@ -128,23 +112,9 @@ async def retrieve_episodes(
     result, _, _ = await driver.execute_query(
         query,
         reference_time=reference_time,
-        source=source.name if source is not None else None,
         num_episodes=last_n,
-        group_ids=group_ids,
+        **query_params,
     )
 
-    episodes = [
-        EpisodicNode(
-            content=record['content'],
-            created_at=parse_db_date(record['created_at'])
-            or datetime.min.replace(tzinfo=timezone.utc),
-            valid_at=parse_db_date(record['valid_at']) or datetime.min.replace(tzinfo=timezone.utc),
-            uuid=record['uuid'],
-            group_id=record['group_id'],
-            source=EpisodeType.from_str(record['source']),
-            name=record['name'],
-            source_description=record['source_description'],
-        )
-        for record in result
-    ]
+    episodes = [get_episodic_node_from_record(record) for record in result]
     return list(reversed(episodes))  # Return in chronological order