graphiti-core 0.18.0__py3-none-any.whl → 0.18.2__py3-none-any.whl

graphiti_core/nodes.py

@@ -25,29 +25,22 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
 from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.nodes.node_db_queries import (
-    COMMUNITY_NODE_SAVE,
-    ENTITY_NODE_SAVE,
+    COMMUNITY_NODE_RETURN,
+    ENTITY_NODE_RETURN,
+    EPISODIC_NODE_RETURN,
     EPISODIC_NODE_SAVE,
+    get_community_node_save_query,
+    get_entity_node_save_query,
 )
 from graphiti_core.utils.datetime_utils import utc_now
 
 logger = logging.getLogger(__name__)
 
-ENTITY_NODE_RETURN: LiteralString = """
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes"""
-
 
 class EpisodeType(Enum):
     """
@@ -96,18 +89,26 @@ class Node(BaseModel, ABC):
     async def save(self, driver: GraphDriver): ...
 
     async def delete(self, driver: GraphDriver):
-        result = await driver.execute_query(
-            """
-        MATCH (n:Entity|Episodic|Community {uuid: $uuid})
-        DETACH DELETE n
-        """,
-            uuid=self.uuid,
-        )
+        if driver.provider == GraphProvider.FALKORDB:
+            for label in ['Entity', 'Episodic', 'Community']:
+                await driver.execute_query(
+                    f"""
+                    MATCH (n:{label} {{uuid: $uuid}})
+                    DETACH DELETE n
+                    """,
+                    uuid=self.uuid,
+                )
+        else:
+            await driver.execute_query(
+                """
+                MATCH (n:Entity|Episodic|Community {uuid: $uuid})
+                DETACH DELETE n
+                """,
+                uuid=self.uuid,
+            )
 
         logger.debug(f'Deleted Node: {self.uuid}')
 
-        return result
-
     def __hash__(self):
         return hash(self.uuid)
 
@@ -118,15 +119,23 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_group_id(cls, driver: GraphDriver, group_id: str):
-        await driver.execute_query(
-            """
-        MATCH (n:Entity|Episodic|Community {group_id: $group_id})
-        DETACH DELETE n
-        """,
-            group_id=group_id,
-        )
-
-        return 'SUCCESS'
+        if driver.provider == GraphProvider.FALKORDB:
+            for label in ['Entity', 'Episodic', 'Community']:
+                await driver.execute_query(
+                    f"""
+                    MATCH (n:{label} {{group_id: $group_id}})
+                    DETACH DELETE n
+                    """,
+                    group_id=group_id,
+                )
+        else:
+            await driver.execute_query(
+                """
+                MATCH (n:Entity|Episodic|Community {group_id: $group_id})
+                DETACH DELETE n
+                """,
+                group_id=group_id,
+            )
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
@@ -169,17 +178,10 @@ class EpisodicNode(Node):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (e:Episodic {uuid: $uuid})
-            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
-        """,
+            MATCH (e:Episodic {uuid: $uuid})
+            RETURN
+            """
+            + EPISODIC_NODE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -195,18 +197,11 @@ class EpisodicNode(Node):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (e:Episodic) WHERE e.uuid IN $uuids
+            MATCH (e:Episodic)
+            WHERE e.uuid IN $uuids
             RETURN DISTINCT
-            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
-        """,
+            """
+            + EPISODIC_NODE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -228,22 +223,17 @@ class EpisodicNode(Node):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (e:Episodic) WHERE e.group_id IN $group_ids
-        """
+            MATCH (e:Episodic)
+            WHERE e.group_id IN $group_ids
+            """
             + cursor_query
             + """
             RETURN DISTINCT
-            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 e.uuid DESC
-        """
+            """
+            + EPISODIC_NODE_RETURN
+            + """
+            ORDER BY uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -259,18 +249,10 @@ class EpisodicNode(Node):
     async def get_by_entity_node_uuid(cls, driver: GraphDriver, entity_node_uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (e:Episodic)-[r:MENTIONS]->(n:Entity {uuid: $entity_node_uuid})
+            MATCH (e:Episodic)-[r:MENTIONS]->(n:Entity {uuid: $entity_node_uuid})
             RETURN DISTINCT
-            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
-        """,
+            """
+            + EPISODIC_NODE_RETURN,
             entity_node_uuid=entity_node_uuid,
             routing_='r',
         )
@@ -297,11 +279,14 @@ class EntityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
-        query: LiteralString = """
+        records, _, _ = await driver.execute_query(
+            """
             MATCH (n:Entity {uuid: $uuid})
             RETURN n.name_embedding AS name_embedding
-        """
-        records, _, _ = await driver.execute_query(query, uuid=self.uuid, routing_='r')
+            """,
+            uuid=self.uuid,
+            routing_='r',
+        )
 
         if len(records) == 0:
             raise NodeNotFoundError(self.uuid)
@@ -317,12 +302,12 @@ class EntityNode(Node):
             'summary': self.summary,
             'created_at': self.created_at,
         }
-
         entity_data.update(self.attributes or {})
 
+        labels = ':'.join(self.labels + ['Entity'])
+
         result = await driver.execute_query(
-            ENTITY_NODE_SAVE,
-            labels=self.labels + ['Entity'],
+            get_entity_node_save_query(driver.provider, labels),
             entity_data=entity_data,
         )
 
@@ -332,14 +317,12 @@ class EntityNode(Node):
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
-        query = (
-            """
-                                                                                MATCH (n:Entity {uuid: $uuid})
-                                                                                """
-            + ENTITY_NODE_RETURN
-        )
         records, _, _ = await driver.execute_query(
-            query,
+            """
+            MATCH (n:Entity {uuid: $uuid})
+            RETURN
+            """
+            + ENTITY_NODE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -355,8 +338,10 @@ class EntityNode(Node):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Entity) WHERE n.uuid IN $uuids
-        """
+            MATCH (n:Entity)
+            WHERE n.uuid IN $uuids
+            RETURN
+            """
             + ENTITY_NODE_RETURN,
             uuids=uuids,
             routing_='r',
@@ -379,22 +364,26 @@ class EntityNode(Node):
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
         with_embeddings_query: LiteralString = (
             """,
-                n.name_embedding AS name_embedding
-                """
+            n.name_embedding AS name_embedding
+            """
             if with_embeddings
             else ''
         )
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Entity) WHERE n.group_id IN $group_ids
-        """
+            MATCH (n:Entity)
+            WHERE n.group_id IN $group_ids
+            """
             + cursor_query
+            + """
+            RETURN
+            """
             + ENTITY_NODE_RETURN
             + with_embeddings_query
             + """
-        ORDER BY n.uuid DESC
-        """
+            ORDER BY n.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -413,7 +402,7 @@ class CommunityNode(Node):
 
     async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
-            COMMUNITY_NODE_SAVE,
+            get_community_node_save_query(driver.provider),
             uuid=self.uuid,
             name=self.name,
             group_id=self.group_id,
@@ -436,11 +425,14 @@ class CommunityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
-        query: LiteralString = """
+        records, _, _ = await driver.execute_query(
+            """
             MATCH (c:Community {uuid: $uuid})
             RETURN c.name_embedding AS name_embedding
-        """
-        records, _, _ = await driver.execute_query(query, uuid=self.uuid, routing_='r')
+            """,
+            uuid=self.uuid,
+            routing_='r',
+        )
 
         if len(records) == 0:
             raise NodeNotFoundError(self.uuid)
@@ -451,14 +443,10 @@ class CommunityNode(Node):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community {uuid: $uuid})
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary
-        """,
+            MATCH (n:Community {uuid: $uuid})
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -474,14 +462,11 @@ 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
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary
-        """,
+            MATCH (n:Community)
+            WHERE n.uuid IN $uuids
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -503,18 +488,17 @@ class CommunityNode(Node):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community) WHERE n.group_id IN $group_ids
-        """
+            MATCH (n:Community)
+            WHERE n.group_id IN $group_ids
+            """
             + cursor_query
             + """
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary
-        ORDER BY n.uuid DESC
-        """
+            RETURN
+            """
+            + COMMUNITY_NODE_RETURN
+            + """
+            ORDER BY n.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -586,6 +570,7 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
 async def create_entity_node_embeddings(embedder: EmbedderClient, nodes: list[EntityNode]):
     if not nodes:  # Handle empty list case
         return
+
     name_embeddings = await embedder.create_batch([node.name for node in nodes])
     for node, name_embedding in zip(nodes, name_embeddings, strict=True):
         node.name_embedding = name_embedding

graphiti_core/prompts/dedupe_nodes.py

@@ -34,7 +34,7 @@ class NodeDuplicate(BaseModel):
     )
     duplicates: list[int] = Field(
         ...,
-        description='idx of all duplicate entities.',
+        description='idx of all entities that are a duplicate of the entity with the above id.',
     )
 
 

graphiti_core/prompts/extract_edges.py

@@ -68,6 +68,10 @@ def edge(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
+<FACT TYPES>
+{context['edge_types']}
+</FACT TYPES>
+
 <PREVIOUS_MESSAGES>
 {json.dumps([ep for ep in context['previous_episodes']], indent=2)}
 </PREVIOUS_MESSAGES>
@@ -84,10 +88,6 @@ def edge(context: dict[str, Any]) -> list[Message]:
 {context['reference_time']}  # ISO 8601 (UTC); used to resolve relative time mentions
 </REFERENCE_TIME>
 
-<FACT TYPES>
-{context['edge_types']}
-</FACT TYPES>
-
 # TASK
 Extract all factual relationships between the given ENTITIES based on the CURRENT MESSAGE.
 Only extract facts that:

graphiti_core/prompts/extract_nodes.py

@@ -52,6 +52,13 @@ class EntityClassification(BaseModel):
     )
 
 
+class EntitySummary(BaseModel):
+    summary: str = Field(
+        ...,
+        description='Summary containing the important information about the entity. Under 250 words',
+    )
+
+
 class Prompt(Protocol):
     extract_message: PromptVersion
     extract_json: PromptVersion
@@ -59,6 +66,7 @@ class Prompt(Protocol):
     reflexion: PromptVersion
     classify_nodes: PromptVersion
     extract_attributes: PromptVersion
+    extract_summary: PromptVersion
 
 
 class Versions(TypedDict):
@@ -68,6 +76,7 @@ class Versions(TypedDict):
     reflexion: PromptFunction
     classify_nodes: PromptFunction
     extract_attributes: PromptFunction
+    extract_summary: PromptFunction
 
 
 def extract_message(context: dict[str, Any]) -> list[Message]:
@@ -75,6 +84,10 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
     Your primary task is to extract and classify the speaker and other significant entities mentioned in the conversation."""
 
     user_prompt = f"""
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
+
 <PREVIOUS MESSAGES>
 {json.dumps([ep for ep in context['previous_episodes']], indent=2)}
 </PREVIOUS MESSAGES>
@@ -83,10 +96,6 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
 {context['episode_content']}
 </CURRENT MESSAGE>
 
-<ENTITY TYPES>
-{context['entity_types']}
-</ENTITY TYPES>
-
 Instructions:
 
 You are given a conversation context and a CURRENT MESSAGE. Your task is to extract **entity nodes** mentioned **explicitly or implicitly** in the CURRENT MESSAGE.
@@ -124,15 +133,16 @@ def extract_json(context: dict[str, Any]) -> list[Message]:
     Your primary task is to extract and classify relevant entities from JSON files"""
 
     user_prompt = f"""
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
+
 <SOURCE DESCRIPTION>:
 {context['source_description']}
 </SOURCE DESCRIPTION>
 <JSON>
 {context['episode_content']}
 </JSON>
-<ENTITY TYPES>
-{context['entity_types']}
-</ENTITY TYPES>
 
 {context['custom_prompt']}
 
@@ -155,13 +165,14 @@ def extract_text(context: dict[str, Any]) -> list[Message]:
     Your primary task is to extract and classify the speaker and other significant entities mentioned in the provided text."""
 
     user_prompt = f"""
-<TEXT>
-{context['episode_content']}
-</TEXT>
 <ENTITY TYPES>
 {context['entity_types']}
 </ENTITY TYPES>
 
+<TEXT>
+{context['episode_content']}
+</TEXT>
+
 Given the above text, extract entities from the TEXT that are explicitly or implicitly mentioned.
 For each entity extracted, also determine its entity type based on the provided ENTITY TYPES and their descriptions.
 Indicate the classified entity type by providing its entity_type_id.
@@ -257,9 +268,39 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
         Guidelines:
         1. Do not hallucinate entity property values if they cannot be found in the current context.
         2. Only use the provided MESSAGES and ENTITY to set attribute values.
+        
+        <ENTITY>
+        {context['node']}
+        </ENTITY>
+        """,
+        ),
+    ]
+
+
+def extract_summary(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that extracts entity summaries from the provided text.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+
+        <MESSAGES>
+        {json.dumps(context['previous_episodes'], indent=2)}
+        {json.dumps(context['episode_content'], indent=2)}
+        </MESSAGES>
+
+        Given the above MESSAGES and the following ENTITY, update the summary that combines relevant information about the entity
+        from the messages and relevant information from the existing summary.
+        
+        Guidelines:
+        1. Do not hallucinate entity summary information if they cannot be found in the current context.
+        2. Only use the provided MESSAGES and ENTITY to set attribute values.
         3. The summary attribute represents a summary of the ENTITY, and should be updated with new information about the Entity from the MESSAGES. 
             Summaries must be no longer than 250 words.
-        
+
         <ENTITY>
         {context['node']}
         </ENTITY>
@@ -273,6 +314,7 @@ versions: Versions = {
     'extract_json': extract_json,
     'extract_text': extract_text,
     'reflexion': reflexion,
+    'extract_summary': extract_summary,
     'classify_nodes': classify_nodes,
     'extract_attributes': extract_attributes,
 }

graphiti_core/search/search_filters.py

@@ -72,7 +72,7 @@ def edge_search_filter_query_constructor(
 
     if filters.edge_types is not None:
         edge_types = filters.edge_types
-        edge_types_filter = '\nAND r.name in $edge_types'
+        edge_types_filter = '\nAND e.name in $edge_types'
         filter_query += edge_types_filter
         filter_params['edge_types'] = edge_types
 
@@ -88,7 +88,7 @@ def edge_search_filter_query_constructor(
                 filter_params['valid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
+                '(e.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -113,7 +113,7 @@ def edge_search_filter_query_constructor(
                 filter_params['invalid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
+                '(e.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -138,7 +138,7 @@ def edge_search_filter_query_constructor(
                 filter_params['created_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
+                '(e.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
@@ -163,7 +163,7 @@ def edge_search_filter_query_constructor(
                 filter_params['expired_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
+                '(e.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''