graphiti-core 0.12.1__tar.gz → 0.12.2__tar.gz

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.12.1
+Version: 0.12.2
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/graphiti.py

@@ -63,6 +63,7 @@ from graphiti_core.utils.maintenance.community_operations import (
     update_community,
 )
 from graphiti_core.utils.maintenance.edge_operations import (
+    build_duplicate_of_edges,
     build_episodic_edges,
     extract_edges,
     resolve_extracted_edge,
@@ -375,7 +376,7 @@ class Graphiti:
             )
 
             # Extract edges and resolve nodes
-            (nodes, uuid_map), extracted_edges = await semaphore_gather(
+            (nodes, uuid_map, node_duplicates), extracted_edges = await semaphore_gather(
                 resolve_extracted_nodes(
                     self.clients,
                     extracted_nodes,
@@ -384,7 +385,13 @@ class Graphiti:
                     entity_types,
                 ),
                 extract_edges(
-                    self.clients, episode, extracted_nodes, previous_episodes, group_id, edge_types
+                    self.clients,
+                    episode,
+                    extracted_nodes,
+                    previous_episodes,
+                    edge_type_map or edge_type_map_default,
+                    group_id,
+                    edge_types,
                 ),
             )
 
@@ -404,7 +411,9 @@ class Graphiti:
                 ),
             )
 
-            entity_edges = resolved_edges + invalidated_edges
+            duplicate_of_edges = build_duplicate_of_edges(episode, now, node_duplicates)
+
+            entity_edges = resolved_edges + invalidated_edges + duplicate_of_edges
 
             episodic_edges = build_episodic_edges(nodes, episode, now)
 
@@ -691,7 +700,7 @@ class Graphiti:
         if edge.fact_embedding is None:
             await edge.generate_embedding(self.embedder)
 
-        resolved_nodes, uuid_map = await resolve_extracted_nodes(
+        resolved_nodes, uuid_map, _ = await resolve_extracted_nodes(
             self.clients,
             [source_node, target_node],
         )

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/prompts/dedupe_nodes.py

@@ -26,12 +26,16 @@ class NodeDuplicate(BaseModel):
     id: int = Field(..., description='integer id of the entity')
     duplicate_idx: int = Field(
         ...,
-        description='idx of the duplicate node. If no duplicate nodes are found, default to -1.',
+        description='idx of the duplicate entity. If no duplicate entities are found, default to -1.',
     )
     name: str = Field(
         ...,
         description='Name of the entity. Should be the most complete and descriptive name possible.',
     )
+    additional_duplicates: list[int] = Field(
+        ...,
+        description='idx of additional duplicate entities. Use this list if the entity has multiple duplicates among existing entities.',
+    )
 
 
 class NodeResolutions(BaseModel):

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/prompts/extract_edges.py

@@ -97,6 +97,8 @@ Only extract facts that:
 - 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
+- The FACT TYPES each contain their fact_type_signature which represents the entity types which that fact_type is defined for.
+    A Type of Entity in the signature represents any extracted entity (it is a generic universal type for all entities).
 
 You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
 

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/prompts/extract_nodes.py

@@ -90,6 +90,8 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
 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.
+Pronoun references such as he/she/they or this/that/those should be disambiguated to the names of the 
+reference entities.
 
 1. **Speaker Extraction**: Always extract the speaker (the part before the colon `:` in each dialogue line) as the first entity node.
    - If the speaker is mentioned again in the message, treat both mentions as a **single entity**.

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/utils/bulk_utils.py

@@ -198,6 +198,7 @@ async def extract_nodes_and_edges_bulk(
                 episode,
                 extracted_nodes_bulk[i],
                 previous_episodes_list[i],
+                {},
                 episode.group_id,
             )
             for i, episode in enumerate(episodes)

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/utils/maintenance/edge_operations.py

@@ -19,7 +19,9 @@ from datetime import datetime
 from time import time
 
 from pydantic import BaseModel
+from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import (
     CommunityEdge,
     EntityEdge,
@@ -27,7 +29,7 @@ from graphiti_core.edges import (
     create_entity_edge_embeddings,
 )
 from graphiti_core.graphiti_types import GraphitiClients
-from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
+from graphiti_core.helpers import DEFAULT_DATABASE, MAX_REFLEXION_ITERATIONS, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.llm_client.config import ModelSize
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
@@ -61,6 +63,28 @@ def build_episodic_edges(
     return episodic_edges
 
 
+def build_duplicate_of_edges(
+    episode: EpisodicNode,
+    created_at: datetime,
+    duplicate_nodes: list[tuple[EntityNode, EntityNode]],
+) -> list[EntityEdge]:
+    is_duplicate_of_edges: list[EntityEdge] = [
+        EntityEdge(
+            source_node_uuid=source_node.uuid,
+            target_node_uuid=target_node.uuid,
+            name='IS_DUPLICATE_OF',
+            group_id=episode.group_id,
+            fact=f'{source_node.name} is a duplicate of {target_node.name}',
+            episodes=[episode.uuid],
+            created_at=created_at,
+            valid_at=created_at,
+        )
+        for source_node, target_node in duplicate_nodes
+    ]
+
+    return is_duplicate_of_edges
+
+
 def build_community_edges(
     entity_nodes: list[EntityNode],
     community_node: CommunityNode,
@@ -84,6 +108,7 @@ async def extract_edges(
     episode: EpisodicNode,
     nodes: list[EntityNode],
     previous_episodes: list[EpisodicNode],
+    edge_type_map: dict[tuple[str, str], list[str]],
     group_id: str = '',
     edge_types: dict[str, BaseModel] | None = None,
 ) -> list[EntityEdge]:
@@ -92,10 +117,17 @@ async def extract_edges(
     extract_edges_max_tokens = 16384
     llm_client = clients.llm_client
 
+    edge_type_signature_map: dict[str, tuple[str, str]] = {
+        edge_type: signature
+        for signature, edge_types in edge_type_map.items()
+        for edge_type in edge_types
+    }
+
     edge_types_context = (
         [
             {
                 'fact_type_name': type_name,
+                'fact_type_signature': edge_type_signature_map.get(type_name, ('Entity', 'Entity')),
                 'fact_type_description': type_model.__doc__,
             }
             for type_name, type_model in edge_types.items()
@@ -107,7 +139,10 @@ async def extract_edges(
     # Prepare context for LLM
     context = {
         'episode_content': episode.content,
-        'nodes': [{'id': idx, 'name': node.name} for idx, node in enumerate(nodes)],
+        'nodes': [
+            {'id': idx, 'name': node.name, 'entity_types': node.labels}
+            for idx, node in enumerate(nodes)
+        ],
         'previous_episodes': [ep.content for ep in previous_episodes],
         'reference_time': episode.valid_at,
         'edge_types': edge_types_context,
@@ -570,3 +605,34 @@ async def dedupe_edge_list(
         unique_edges.append(edge)
 
     return unique_edges
+
+
+async def filter_existing_duplicate_of_edges(
+    driver: GraphDriver, duplicates_node_tuples: list[tuple[EntityNode, EntityNode]]
+) -> list[tuple[EntityNode, EntityNode]]:
+    query: LiteralString = """
+        UNWIND $duplicate_node_uuids AS duplicate_tuple
+        MATCH (n:Entity {uuid: duplicate_tuple[0]})-[r:RELATES_TO {name: 'IS_DUPLICATE_OF'}]->(m:Entity {uuid: duplicate_tuple[1]})
+        RETURN DISTINCT
+            n.uuid AS source_uuid,
+            m.uuid AS target_uuid
+    """
+
+    duplicate_nodes_map = {
+        (source.uuid, target.uuid): (source, target) for source, target in duplicates_node_tuples
+    }
+
+    records, _, _ = await driver.execute_query(
+        query,
+        duplicate_node_uuids=list(duplicate_nodes_map.keys()),
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+
+    # Remove duplicates that already have the IS_DUPLICATE_OF edge
+    for record in records:
+        duplicate_tuple = (record.get('source_uuid'), record.get('target_uuid'))
+        if duplicate_nodes_map.get(duplicate_tuple):
+            duplicate_nodes_map.pop(duplicate_tuple)
+
+    return list(duplicate_nodes_map.values())

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/graphiti_core/utils/maintenance/node_operations.py

@@ -40,6 +40,7 @@ from graphiti_core.search.search_config import SearchResults
 from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.utils.datetime_utils import utc_now
+from graphiti_core.utils.maintenance.edge_operations import filter_existing_duplicate_of_edges
 
 logger = logging.getLogger(__name__)
 
@@ -225,8 +226,9 @@ async def resolve_extracted_nodes(
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
     entity_types: dict[str, BaseModel] | None = None,
-) -> tuple[list[EntityNode], dict[str, str]]:
+) -> tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]:
     llm_client = clients.llm_client
+    driver = clients.driver
 
     search_results: list[SearchResults] = await semaphore_gather(
         *[
@@ -295,9 +297,10 @@ async def resolve_extracted_nodes(
 
     resolved_nodes: list[EntityNode] = []
     uuid_map: dict[str, str] = {}
+    node_duplicates: list[tuple[EntityNode, EntityNode]] = []
     for resolution in node_resolutions:
-        resolution_id = resolution.get('id', -1)
-        duplicate_idx = resolution.get('duplicate_idx', -1)
+        resolution_id: int = resolution.get('id', -1)
+        duplicate_idx: int = resolution.get('duplicate_idx', -1)
 
         extracted_node = extracted_nodes[resolution_id]
 
@@ -312,9 +315,21 @@ async def resolve_extracted_nodes(
         resolved_nodes.append(resolved_node)
         uuid_map[extracted_node.uuid] = resolved_node.uuid
 
+        additional_duplicates: list[int] = resolution.get('additional_duplicates', [])
+        for idx in additional_duplicates:
+            existing_node = existing_nodes[idx] if idx < len(existing_nodes) else resolved_node
+            if existing_node == resolved_node:
+                continue
+
+            node_duplicates.append((resolved_node, existing_nodes[idx]))
+
     logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
 
-    return resolved_nodes, uuid_map
+    new_node_duplicates: list[
+        tuple[EntityNode, EntityNode]
+    ] = await filter_existing_duplicate_of_edges(driver, node_duplicates)
+
+    return resolved_nodes, uuid_map, new_node_duplicates
 
 
 async def extract_attributes_from_nodes(

{graphiti_core-0.12.1 → graphiti_core-0.12.2}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.12.1"
+version = "0.12.2"
 authors = [
     { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" },
     { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" },