graphiti-core 0.11.6rc9__py3-none-any.whl → 0.12.0__py3-none-any.whl

graphiti_core/nodes.py

@@ -22,13 +22,13 @@ from time import time
 from typing import Any
 from uuid import uuid4
 
-from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_SAVE,
     ENTITY_NODE_SAVE,
@@ -94,9 +94,9 @@ class Node(BaseModel, ABC):
     created_at: datetime = Field(default_factory=lambda: utc_now())
 
     @abstractmethod
-    async def save(self, driver: AsyncDriver): ...
+    async def save(self, driver: GraphDriver): ...
 
-    async def delete(self, driver: AsyncDriver):
+    async def delete(self, driver: GraphDriver):
         result = await driver.execute_query(
             """
         MATCH (n:Entity|Episodic|Community {uuid: $uuid})
@@ -119,7 +119,7 @@ class Node(BaseModel, ABC):
         return False
 
     @classmethod
-    async def delete_by_group_id(cls, driver: AsyncDriver, group_id: str):
+    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})
@@ -132,10 +132,10 @@ class Node(BaseModel, ABC):
         return 'SUCCESS'
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str): ...
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]): ...
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]): ...
 
 
 class EpisodicNode(Node):
@@ -150,7 +150,7 @@ class EpisodicNode(Node):
         default_factory=list,
     )
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             EPISODIC_NODE_SAVE,
             uuid=self.uuid,
@@ -165,12 +165,12 @@ class EpisodicNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic {uuid: $uuid})
@@ -197,7 +197,7 @@ class EpisodicNode(Node):
         return episodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic) WHERE e.uuid IN $uuids
@@ -224,7 +224,7 @@ class EpisodicNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -263,7 +263,7 @@ class EpisodicNode(Node):
         return episodes
 
     @classmethod
-    async def get_by_entity_node_uuid(cls, driver: AsyncDriver, entity_node_uuid: str):
+    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})
@@ -304,7 +304,7 @@ class EntityNode(Node):
 
         return self.name_embedding
 
-    async def load_name_embedding(self, driver: AsyncDriver):
+    async def load_name_embedding(self, driver: GraphDriver):
         query: LiteralString = """
             MATCH (n:Entity {uuid: $uuid})
             RETURN n.name_embedding AS name_embedding
@@ -318,7 +318,7 @@ class EntityNode(Node):
 
         self.name_embedding = records[0]['name_embedding']
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         entity_data: dict[str, Any] = {
             'uuid': self.uuid,
             'name': self.name,
@@ -337,16 +337,16 @@ class EntityNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         query = (
             """
-                                                        MATCH (n:Entity {uuid: $uuid})
-                                                        """
+                                                                    MATCH (n:Entity {uuid: $uuid})
+                                                                    """
             + ENTITY_NODE_RETURN
         )
         records, _, _ = await driver.execute_query(
@@ -364,7 +364,7 @@ class EntityNode(Node):
         return nodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity) WHERE n.uuid IN $uuids
@@ -382,7 +382,7 @@ class EntityNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -416,7 +416,7 @@ class CommunityNode(Node):
     name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
     summary: str = Field(description='region summary of member nodes', default_factory=str)
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             COMMUNITY_NODE_SAVE,
             uuid=self.uuid,
@@ -428,7 +428,7 @@ class CommunityNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
@@ -441,7 +441,7 @@ class CommunityNode(Node):
 
         return self.name_embedding
 
-    async def load_name_embedding(self, driver: AsyncDriver):
+    async def load_name_embedding(self, driver: GraphDriver):
         query: LiteralString = """
             MATCH (c:Community {uuid: $uuid})
             RETURN c.name_embedding AS name_embedding
@@ -456,7 +456,7 @@ class CommunityNode(Node):
         self.name_embedding = records[0]['name_embedding']
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community {uuid: $uuid})
@@ -480,7 +480,7 @@ class CommunityNode(Node):
         return nodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community) WHERE n.uuid IN $uuids
@@ -503,7 +503,7 @@ class CommunityNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -542,8 +542,8 @@ class CommunityNode(Node):
 def get_episodic_node_from_record(record: Any) -> EpisodicNode:
     return EpisodicNode(
         content=record['content'],
-        created_at=record['created_at'].to_native().timestamp(),
-        valid_at=(record['valid_at'].to_native()),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
+        valid_at=parse_db_date(record['valid_at']),  # type: ignore
         uuid=record['uuid'],
         group_id=record['group_id'],
         source=EpisodeType.from_str(record['source']),
@@ -559,7 +559,7 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         name=record['name'],
         group_id=record['group_id'],
         labels=record['labels'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
         attributes=record['attributes'],
     )
@@ -580,7 +580,7 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
         name=record['name'],
         group_id=record['group_id'],
         name_embedding=record['name_embedding'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
     )
 

graphiti_core/prompts/dedupe_edges.py

@@ -27,6 +27,11 @@ class EdgeDuplicate(BaseModel):
         ...,
         description='id of the duplicate fact. If no duplicate facts are found, default to -1.',
     )
+    contradicted_facts: list[int] = Field(
+        ...,
+        description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
+    )
+    fact_type: str = Field(..., description='One of the provided fact types or DEFAULT')
 
 
 class UniqueFact(BaseModel):
@@ -41,11 +46,13 @@ class UniqueFacts(BaseModel):
 class Prompt(Protocol):
     edge: PromptVersion
     edge_list: PromptVersion
+    resolve_edge: PromptVersion
 
 
 class Versions(TypedDict):
     edge: PromptFunction
     edge_list: PromptFunction
+    resolve_edge: PromptFunction
 
 
 def edge(context: dict[str, Any]) -> list[Message]:
@@ -106,4 +113,48 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'edge': edge, 'edge_list': edge_list}
+def resolve_edge(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that de-duplicates facts from fact lists and determines which existing '
+            'facts are contradicted by the new fact.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        <NEW FACT>
+        {context['new_edge']}
+        </NEW FACT>
+        
+        <EXISTING FACTS>
+        {context['existing_edges']}
+        </EXISTING FACTS>
+        <FACT INVALIDATION CANDIDATES>
+        {context['edge_invalidation_candidates']}
+        </FACT INVALIDATION CANDIDATES>
+        
+        <FACT TYPES>
+        {context['edge_types']}
+        </FACT TYPES>
+        
+
+        Task:
+        If the NEW FACT represents the same factual information as any fact in EXISTING FACTS, return the idx of the duplicate fact.
+        If the NEW FACT is not a duplicate of any of the EXISTING FACTS, return -1.
+        
+        Given the predefined FACT TYPES, determine if the NEW FACT should be classified as one of these types.
+        Return the fact type as fact_type or DEFAULT if NEW FACT is not one of the FACT TYPES.
+        
+        Based on the provided FACT INVALIDATION CANDIDATES and NEW FACT, determine which existing facts the new fact contradicts.
+        Return a list containing all idx's of the facts that are contradicted by the NEW FACT.
+        If there are no contradicted facts, return an empty list.
+
+        Guidelines:
+        1. The facts do not need to be completely identical to be duplicates, they just need to express the same information.
+        """,
+        ),
+    ]
+
+
+versions: Versions = {'edge': edge, 'edge_list': edge_list, 'resolve_edge': resolve_edge}

graphiti_core/prompts/dedupe_nodes.py

@@ -23,21 +23,31 @@ from .models import Message, PromptFunction, PromptVersion
 
 
 class NodeDuplicate(BaseModel):
-    duplicate_node_id: int = Field(
+    id: int = Field(..., description='integer id of the entity')
+    duplicate_idx: int = Field(
         ...,
-        description='id of the duplicate node. If no duplicate nodes are found, default to -1.',
+        description='idx of the duplicate node. If no duplicate nodes are found, default to -1.',
     )
-    name: str = Field(..., description='Name of the entity.')
+    name: str = Field(
+        ...,
+        description='Name of the entity. Should be the most complete and descriptive name possible.',
+    )
+
+
+class NodeResolutions(BaseModel):
+    entity_resolutions: list[NodeDuplicate] = Field(..., description='List of resolved nodes')
 
 
 class Prompt(Protocol):
     node: PromptVersion
     node_list: PromptVersion
+    nodes: PromptVersion
 
 
 class Versions(TypedDict):
     node: PromptFunction
     node_list: PromptFunction
+    nodes: PromptFunction
 
 
 def node(context: dict[str, Any]) -> list[Message]:
@@ -89,6 +99,71 @@ def node(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def nodes(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that determines whether or not ENTITIES extracted from a conversation are duplicates'
+            'of existing entities.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        <PREVIOUS MESSAGES>
+        {json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+        </PREVIOUS MESSAGES>
+        <CURRENT MESSAGE>
+        {context['episode_content']}
+        </CURRENT MESSAGE>
+        
+        
+        Each of the following ENTITIES were extracted from the CURRENT MESSAGE.
+        Each entity in ENTITIES is represented as a JSON object with the following structure:
+        {{
+            id: integer id of the entity,
+            name: "name of the entity",
+            entity_type: "ontological classification of the entity",
+            entity_type_description: "Description of what the entity type represents",
+            duplication_candidates: [
+                {{
+                    idx: integer index of the candidate entity,
+                    name: "name of the candidate entity",
+                    entity_type: "ontological classification of the candidate entity",
+                    ...<additional attributes>
+                }}
+            ]
+        }}
+        
+        <ENTITIES>
+        {json.dumps(context['extracted_nodes'], indent=2)}
+        </ENTITIES>
+        
+        <EXISTING ENTITIES>
+        {json.dumps(context['existing_nodes'], indent=2)}
+        </EXISTING ENTITIES>
+
+        For each of the above ENTITIES, determine if the entity is a duplicate of any of the EXISTING ENTITIES.
+
+        Entities should only be considered duplicates if they refer to the *same real-world object or concept*.
+
+        Do NOT mark entities as duplicates if:
+        - They are related but distinct.
+        - They have similar names or purposes but refer to separate instances or concepts.
+
+        Task:
+        Your response will be a list called entity_resolutions which contains one entry for each entity.
+        
+        For each entity, return the id of the entity as id, the name of the entity as name, and the duplicate_idx
+        as an integer.
+        
+        - If an entity is a duplicate of one of the EXISTING ENTITIES, return the idx of the candidate it is a 
+        duplicate of.
+        - If an entity is not a duplicate of one of the EXISTING ENTITIES, return the -1 as the duplication_idx
+        """,
+        ),
+    ]
+
+
 def node_list(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
@@ -126,4 +201,4 @@ def node_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'node': node, 'node_list': node_list}
+versions: Versions = {'node': node, 'node_list': node_list, 'nodes': nodes}

graphiti_core/prompts/extract_edges.py

@@ -24,8 +24,8 @@ from .models import Message, PromptFunction, PromptVersion
 
 class Edge(BaseModel):
     relation_type: str = Field(..., description='FACT_PREDICATE_IN_SCREAMING_SNAKE_CASE')
-    source_entity_name: str = Field(..., description='The name of the source entity of the fact.')
-    target_entity_name: str = Field(..., description='The name of the target entity of the fact.')
+    source_entity_id: int = Field(..., description='The id of the source entity of the fact.')
+    target_entity_id: int = Field(..., description='The id of the target entity of the fact.')
     fact: str = Field(..., description='')
     valid_at: str | None = Field(
         None,
@@ -48,11 +48,13 @@ class MissingFacts(BaseModel):
 class Prompt(Protocol):
     edge: PromptVersion
     reflexion: PromptVersion
+    extract_attributes: PromptVersion
 
 
 class Versions(TypedDict):
     edge: PromptFunction
     reflexion: PromptFunction
+    extract_attributes: PromptFunction
 
 
 def edge(context: dict[str, Any]) -> list[Message]:
@@ -75,19 +77,26 @@ def edge(context: dict[str, Any]) -> list[Message]:
 </CURRENT_MESSAGE>
 
 <ENTITIES>
-{context['nodes']}  # Each has: id, label (e.g., Person, Org), name, aliases
+{context['nodes']} 
 </ENTITIES>
 
 <REFERENCE_TIME>
 {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:
 - involve two DISTINCT ENTITIES from the ENTITIES list,
 - are clearly stated or unambiguously implied in the CURRENT MESSAGE,
-- and can be represented as edges in a knowledge graph.
+    and can be represented as edges in a knowledge graph.
+- The FACT TYPES provide a list of the most important types of facts, make sure to extract facts of these types
+- The FACT TYPES are not an exhaustive list, extract all facts from the message even if they do not fit into one
+    of the FACT TYPES
 
 You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
 
@@ -145,4 +154,40 @@ determine if any facts haven't been extracted.
     ]
 
 
-versions: Versions = {'edge': edge, 'reflexion': reflexion}
+def extract_attributes(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that extracts fact properties from the provided text.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+
+        <MESSAGE>
+        {json.dumps(context['episode_content'], indent=2)}
+        </MESSAGE>
+        <REFERENCE TIME>
+        {context['reference_time']}
+        </REFERENCE TIME>
+
+        Given the above MESSAGE, its REFERENCE TIME, and the following FACT, update any of its attributes based on the information provided
+        in MESSAGE. Use the provided attribute descriptions to better understand how each attribute should be determined.
+
+        Guidelines:
+        1. Do not hallucinate entity property values if they cannot be found in the current context.
+        2. Only use the provided MESSAGES and FACT to set attribute values.
+
+        <FACT>
+        {context['fact']}
+        </FACT>
+        """,
+        ),
+    ]
+
+
+versions: Versions = {
+    'edge': edge,
+    'reflexion': reflexion,
+    'extract_attributes': extract_attributes,
+}

graphiti_core/prompts/invalidate_edges.py

@@ -24,7 +24,7 @@ from .models import Message, PromptFunction, PromptVersion
 class InvalidatedEdges(BaseModel):
     contradicted_facts: list[int] = Field(
         ...,
-        description='List of ids of facts that be should invalidated. If no facts should be invalidated, the list should be empty.',
+        description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
     )
 
 

graphiti_core/search/search.py

@@ -18,9 +18,8 @@ import logging
 from collections import defaultdict
 from time import time
 
-from neo4j import AsyncDriver
-
 from graphiti_core.cross_encoder.client import CrossEncoderClient
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
@@ -94,7 +93,7 @@ async def search(
     )
 
     # if group_ids is empty, set it to None
-    group_ids = group_ids if group_ids else None
+    group_ids = group_ids if group_ids and group_ids != [''] else None
     edges, nodes, episodes, communities = await semaphore_gather(
         edge_search(
             driver,
@@ -160,7 +159,7 @@ async def search(
 
 
 async def edge_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -174,7 +173,6 @@ async def edge_search(
 ) -> list[EntityEdge]:
     if config is None:
         return []
-
     search_results: list[list[EntityEdge]] = list(
         await semaphore_gather(
             *[
@@ -261,7 +259,7 @@ async def edge_search(
 
 
 async def node_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -275,7 +273,6 @@ async def node_search(
 ) -> list[EntityNode]:
     if config is None:
         return []
-
     search_results: list[list[EntityNode]] = list(
         await semaphore_gather(
             *[
@@ -344,7 +341,7 @@ async def node_search(
 
 
 async def episode_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     _query_vector: list[float],
@@ -356,7 +353,6 @@ async def episode_search(
 ) -> list[EpisodicNode]:
     if config is None:
         return []
-
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
@@ -392,7 +388,7 @@ async def episode_search(
 
 
 async def community_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],