graphiti-core 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

graphiti_core/graphiti.py

@@ -29,6 +29,7 @@ from graphiti_core.llm_client.utils import generate_embedding
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import Reranker, SearchConfig, SearchMethod, hybrid_search
 from graphiti_core.search.search_utils import (
+    RELEVANT_SCHEMA_LIMIT,
     get_relevant_edges,
     get_relevant_nodes,
     hybrid_node_search,
@@ -41,19 +42,23 @@ from graphiti_core.utils.bulk_utils import (
     RawEpisode,
     dedupe_edges_bulk,
     dedupe_nodes_bulk,
+    extract_edge_dates_bulk,
     extract_nodes_and_edges_bulk,
     resolve_edge_pointers,
     retrieve_previous_episodes_bulk,
 )
 from graphiti_core.utils.maintenance.edge_operations import (
-    dedupe_extracted_edges,
     extract_edges,
+    resolve_extracted_edges,
 )
 from graphiti_core.utils.maintenance.graph_data_operations import (
     EPISODE_WINDOW_LEN,
     build_indices_and_constraints,
 )
-from graphiti_core.utils.maintenance.node_operations import dedupe_extracted_nodes, extract_nodes
+from graphiti_core.utils.maintenance.node_operations import (
+    extract_nodes,
+    resolve_extracted_nodes,
+)
 from graphiti_core.utils.maintenance.temporal_operations import (
     extract_edge_dates,
     invalidate_edges,
@@ -175,9 +180,9 @@ class Graphiti:
         await build_indices_and_constraints(self.driver)
 
     async def retrieve_episodes(
-        self,
-        reference_time: datetime,
-        last_n: int = EPISODE_WINDOW_LEN,
+            self,
+            reference_time: datetime,
+            last_n: int = EPISODE_WINDOW_LEN,
     ) -> list[EpisodicNode]:
         """
         Retrieve the last n episodic nodes from the graph.
@@ -205,14 +210,14 @@ class Graphiti:
         return await retrieve_episodes(self.driver, reference_time, last_n)
 
     async def add_episode(
-        self,
-        name: str,
-        episode_body: str,
-        source_description: str,
-        reference_time: datetime,
-        source: EpisodeType = EpisodeType.message,
-        success_callback: Callable | None = None,
-        error_callback: Callable | None = None,
+            self,
+            name: str,
+            episode_body: str,
+            source_description: str,
+            reference_time: datetime,
+            source: EpisodeType = EpisodeType.message,
+            success_callback: Callable | None = None,
+            error_callback: Callable | None = None,
     ):
         """
         Process an episode and update the graph.
@@ -263,7 +268,6 @@ class Graphiti:
 
             nodes: list[EntityNode] = []
             entity_edges: list[EntityEdge] = []
-            episodic_edges: list[EpisodicEdge] = []
             embedder = self.llm_client.get_embedder()
             now = datetime.now()
 
@@ -278,6 +282,8 @@ class Graphiti:
                 valid_at=reference_time,
             )
 
+            # Extract entities as nodes
+
             extracted_nodes = await extract_nodes(self.llm_client, episode, previous_episodes)
             logger.info(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
@@ -286,59 +292,82 @@ class Graphiti:
             await asyncio.gather(
                 *[node.generate_name_embedding(embedder) for node in extracted_nodes]
             )
-            existing_nodes = await get_relevant_nodes(extracted_nodes, self.driver)
+
+            # Resolve extracted nodes with nodes already in the graph
+            existing_nodes_lists: list[list[EntityNode]] = list(
+                await asyncio.gather(
+                    *[get_relevant_nodes([node], self.driver) for node in extracted_nodes]
+                )
+            )
+
             logger.info(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
-            touched_nodes, _, brand_new_nodes = await dedupe_extracted_nodes(
-                self.llm_client, extracted_nodes, existing_nodes
+
+            mentioned_nodes, _ = await resolve_extracted_nodes(
+                self.llm_client, extracted_nodes, existing_nodes_lists
             )
-            logger.info(f'Adjusted touched nodes: {[(n.name, n.uuid) for n in touched_nodes]}')
-            nodes.extend(touched_nodes)
+            logger.info(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
+            nodes.extend(mentioned_nodes)
 
+            # Extract facts as edges given entity nodes
             extracted_edges = await extract_edges(
-                self.llm_client, episode, touched_nodes, previous_episodes
+                self.llm_client, episode, mentioned_nodes, previous_episodes
             )
 
+            # calculate embeddings
             await asyncio.gather(*[edge.generate_embedding(embedder) for edge in extracted_edges])
 
-            existing_edges = await get_relevant_edges(extracted_edges, self.driver)
-            logger.info(f'Existing edges: {[(e.name, e.uuid) for e in existing_edges]}')
+            # Resolve extracted edges with edges already in the graph
+            existing_edges_list: list[list[EntityEdge]] = list(
+                await asyncio.gather(
+                    *[
+                        get_relevant_edges(
+                            [edge],
+                            self.driver,
+                            RELEVANT_SCHEMA_LIMIT,
+                            edge.source_node_uuid,
+                            edge.target_node_uuid,
+                        )
+                        for edge in extracted_edges
+                    ]
+                )
+            )
+            logger.info(
+                f'Existing edges lists: {[(e.name, e.uuid) for edges_lst in existing_edges_list for e in edges_lst]}'
+            )
             logger.info(f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges]}')
 
-            deduped_edges = await dedupe_extracted_edges(
-                self.llm_client,
-                extracted_edges,
-                existing_edges,
+            deduped_edges: list[EntityEdge] = await resolve_extracted_edges(
+                self.llm_client, extracted_edges, existing_edges_list
             )
 
-            edge_touched_node_uuids = [n.uuid for n in brand_new_nodes]
-            for edge in deduped_edges:
-                edge_touched_node_uuids.append(edge.source_node_uuid)
-                edge_touched_node_uuids.append(edge.target_node_uuid)
-
-            for edge in deduped_edges:
-                valid_at, invalid_at, _ = await extract_edge_dates(
-                    self.llm_client,
-                    edge,
-                    episode.valid_at,
-                    episode,
-                    previous_episodes,
-                )
-                edge.valid_at = valid_at
-                edge.invalid_at = invalid_at
-                if edge.invalid_at:
-                    edge.expired_at = datetime.now()
-            for edge in existing_edges:
-                valid_at, invalid_at, _ = await extract_edge_dates(
-                    self.llm_client,
-                    edge,
-                    episode.valid_at,
-                    episode,
-                    previous_episodes,
-                )
+            # Extract dates for the newly extracted edges
+            edge_dates = await asyncio.gather(
+                *[
+                    extract_edge_dates(
+                        self.llm_client,
+                        edge,
+                        episode,
+                        previous_episodes,
+                    )
+                    for edge in deduped_edges
+                ]
+            )
+
+            for i, edge in enumerate(deduped_edges):
+                valid_at = edge_dates[i][0]
+                invalid_at = edge_dates[i][1]
+
                 edge.valid_at = valid_at
                 edge.invalid_at = invalid_at
-                if edge.invalid_at:
-                    edge.expired_at = datetime.now()
+                if edge.invalid_at is not None:
+                    edge.expired_at = now
+
+            entity_edges.extend(deduped_edges)
+
+            existing_edges: list[EntityEdge] = [
+                e for edge_lst in existing_edges_list for e in edge_lst
+            ]
+
             (
                 old_edges_with_nodes_pending_invalidation,
                 new_edges_with_nodes,
@@ -361,30 +390,18 @@ class Graphiti:
                 for deduped_edge in deduped_edges:
                     if deduped_edge.uuid == edge.uuid:
                         deduped_edge.expired_at = edge.expired_at
-                edge_touched_node_uuids.append(edge.source_node_uuid)
-                edge_touched_node_uuids.append(edge.target_node_uuid)
             logger.info(f'Invalidated edges: {[(e.name, e.uuid) for e in invalidated_edges]}')
 
-            edges_to_save = existing_edges + deduped_edges
-
-            entity_edges.extend(edges_to_save)
-
-            edge_touched_node_uuids = list(set(edge_touched_node_uuids))
-            involved_nodes = [node for node in nodes if node.uuid in edge_touched_node_uuids]
-
-            logger.info(f'Edge touched nodes: {[(n.name, n.uuid) for n in involved_nodes]}')
+            entity_edges.extend(existing_edges)
 
             logger.info(f'Deduped edges: {[(e.name, e.uuid) for e in deduped_edges]}')
 
-            episodic_edges.extend(
-                build_episodic_edges(
-                    # There may be an overlap between new_nodes and affected_nodes, so we're deduplicating them
-                    involved_nodes,
-                    episode,
-                    now,
-                )
+            episodic_edges: list[EpisodicEdge] = build_episodic_edges(
+                mentioned_nodes,
+                episode,
+                now,
             )
-            # Important to append the episode to the nodes at the end so that self referencing episodic edges are not built
+
             logger.info(f'Built episodic edges: {episodic_edges}')
 
             # Future optimization would be using batch operations to save nodes and edges
@@ -395,9 +412,7 @@ class Graphiti:
 
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
-            # for node in nodes:
-            #     if isinstance(node, EntityNode):
-            #         await node.update_summary(self.driver)
+
             if success_callback:
                 await success_callback(episode)
         except Exception as e:
@@ -407,8 +422,8 @@ class Graphiti:
                 raise e
 
     async def add_episode_bulk(
-        self,
-        bulk_episodes: list[RawEpisode],
+            self,
+            bulk_episodes: list[RawEpisode],
     ):
         """
         Process multiple episodes in bulk and update the graph.
@@ -481,15 +496,18 @@ class Graphiti:
                 *[edge.generate_embedding(embedder) for edge in extracted_edges],
             )
 
-            # Dedupe extracted nodes
-            nodes, uuid_map = await dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes)
+            # Dedupe extracted nodes, compress extracted edges
+            (nodes, uuid_map), extracted_edges_timestamped = await asyncio.gather(
+                dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes),
+                extract_edge_dates_bulk(self.llm_client, extracted_edges, episode_pairs),
+            )
 
             # save nodes to KG
             await asyncio.gather(*[node.save(self.driver) for node in nodes])
 
             # re-map edge pointers so that they don't point to discard dupe nodes
             extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
-                extracted_edges, uuid_map
+                extracted_edges_timestamped, uuid_map
             )
             episodic_edges_with_resolved_pointers: list[EpisodicEdge] = resolve_edge_pointers(
                 episodic_edges, uuid_map
@@ -569,17 +587,19 @@ class Graphiti:
         return edges
 
     async def _search(
-        self,
-        query: str,
-        timestamp: datetime,
-        config: SearchConfig,
-        center_node_uuid: str | None = None,
+            self,
+            query: str,
+            timestamp: datetime,
+            config: SearchConfig,
+            center_node_uuid: str | None = None,
     ):
         return await hybrid_search(
             self.driver, self.llm_client.get_embedder(), query, timestamp, config, center_node_uuid
         )
 
-    async def get_nodes_by_query(self, query: str, limit: int | None = None) -> list[EntityNode]:
+    async def get_nodes_by_query(
+            self, query: str, limit: int = RELEVANT_SCHEMA_LIMIT
+    ) -> list[EntityNode]:
         """
         Retrieve nodes from the graph database based on a text query.
 

graphiti_core/llm_client/openai_client.py

@@ -60,6 +60,5 @@ class OpenAIClient(LLMClient):
             result = response.choices[0].message.content or ''
             return json.loads(result)
         except Exception as e:
-            print(openai_messages)
             logger.error(f'Error in generating LLM response: {e}')
             raise

graphiti_core/prompts/dedupe_edges.py

@@ -23,12 +23,14 @@ from .models import Message, PromptFunction, PromptVersion
 class Prompt(Protocol):
     v1: PromptVersion
     v2: PromptVersion
+    v3: PromptVersion
     edge_list: PromptVersion
 
 
 class Versions(TypedDict):
     v1: PromptFunction
     v2: PromptFunction
+    v3: PromptFunction
     edge_list: PromptFunction
 
 
@@ -41,17 +43,17 @@ def v1(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
-        Given the following context, deduplicate facts from a list of new facts given a list of existing facts:
+        Given the following context, deduplicate facts from a list of new facts given a list of existing edges:
 
-        Existing Facts:
+        Existing Edges:
         {json.dumps(context['existing_edges'], indent=2)}
 
-        New Facts:
+        New Edges:
         {json.dumps(context['extracted_edges'], indent=2)}
 
         Task:
-        If any facts in New Facts is a duplicate of a fact in Existing Facts, 
-        do not return it in the list of unique facts.
+        If any edge in New Edges is a duplicate of an edge in Existing Edges, add their uuids to the output list.
+        When finding duplicates edges, synthesize their facts into a short new fact.
 
         Guidelines:
         1. identical or near identical facts are duplicates
@@ -60,9 +62,11 @@ def v1(context: dict[str, Any]) -> list[Message]:
 
         Respond with a JSON object in the following format:
         {{
-            "unique_facts": [
+            "duplicates": [
                 {{
-                    "uuid": "unique identifier of the fact"
+                    "uuid": "uuid of the new node like 5d643020624c42fa9de13f97b1b3fa39",
+                    "duplicate_of": "uuid of the existing node",
+                    "fact": "one sentence description of the fact"
                 }}
             ]
         }}
@@ -113,6 +117,40 @@ def v2(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def v3(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that de-duplicates edges from edge lists.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        Given the following context, determine whether the New Edge represents any of the edges in the list of Existing Edges.
+
+        Existing Edges:
+        {json.dumps(context['existing_edges'], indent=2)}
+
+        New Edge:
+        {json.dumps(context['extracted_edges'], indent=2)}
+        Task:
+        1. If the New Edges represents the same factual information as any edge in Existing Edges, return 'is_duplicate: true' in the 
+            response. Otherwise, return 'is_duplicate: false'
+        2. If is_duplicate is true, also return the uuid of the existing edge in the response
+
+        Guidelines:
+        1. The facts do not need to be completely identical to be duplicates, they just need to express the same information.
+
+        Respond with a JSON object in the following format:
+            {{
+                "is_duplicate": true or false,
+                "uuid": uuid of the existing edge like "5d643020624c42fa9de13f97b1b3fa39" or null,
+            }}
+        """,
+        ),
+    ]
+
+
 def edge_list(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
@@ -151,4 +189,4 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'v1': v1, 'v2': v2, 'edge_list': edge_list}
+versions: Versions = {'v1': v1, 'v2': v2, 'v3': v3, 'edge_list': edge_list}

graphiti_core/prompts/dedupe_nodes.py

@@ -23,13 +23,15 @@ from .models import Message, PromptFunction, PromptVersion
 class Prompt(Protocol):
     v1: PromptVersion
     v2: PromptVersion
+    v3: PromptVersion
     node_list: PromptVersion
 
 
 class Versions(TypedDict):
     v1: PromptFunction
     v2: PromptFunction
-    node_list: PromptVersion
+    v3: PromptFunction
+    node_list: PromptFunction
 
 
 def v1(context: dict[str, Any]) -> list[Message]:
@@ -53,7 +55,9 @@ def v1(context: dict[str, Any]) -> list[Message]:
         1. start with the list of nodes from New Nodes
         2. If any node in New Nodes is a duplicate of a node in Existing Nodes, replace the new node with the existing
             node in the list
-        3. Respond with the resulting list of nodes
+        3. when deduplicating nodes, synthesize their summaries into a short new summary that contains the relevant information
+            of the summaries of the new and existing nodes
+        4. Respond with the resulting list of nodes
 
         Guidelines:
         1. Use both the name and summary of nodes to determine if they are duplicates, 
@@ -64,6 +68,7 @@ def v1(context: dict[str, Any]) -> list[Message]:
             "new_nodes": [
                 {{
                     "name": "Unique identifier for the node",
+                    "summary": "Brief summary of the node's role or significance"
                 }}
             ]
         }}
@@ -91,20 +96,23 @@ def v2(context: dict[str, Any]) -> list[Message]:
         Important:
         If a node in the new nodes is describing the same entity as a node in the existing nodes, mark it as a duplicate!!!
         Task:
-        If any node in New Nodes is a duplicate of a node in Existing Nodes, add their names to the output list
+        If any node in New Nodes is a duplicate of a node in Existing Nodes, add their uuids to the output list
+        When finding duplicates nodes, synthesize their summaries into a short new summary that contains the 
+        relevant information of the summaries of the new and existing nodes.
 
         Guidelines:
         1. Use both the name and summary of nodes to determine if they are duplicates, 
             duplicate nodes may have different names
-        2. In the output, name should always be the name of the New Node that is a duplicate. duplicate_of should be
-            the name of the Existing Node.
+        2. In the output, uuid should always be the uuid of the New Node that is a duplicate. duplicate_of should be
+            the uuid of the Existing Node.
 
         Respond with a JSON object in the following format:
         {{
             "duplicates": [
                 {{
-                    "name": "name of the new node",
-                    "duplicate_of": "name of the existing node"
+                    "uuid": "uuid of the new node like 5d643020624c42fa9de13f97b1b3fa39",
+                    "duplicate_of": "uuid of the existing node",
+                    "summary": "Brief summary of the node's role or significance. Takes information from the new and existing nodes"
                 }}
             ]
         }}
@@ -113,6 +121,44 @@ def v2(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def v3(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that de-duplicates nodes from node lists.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        Given the following context, determine whether the New Node represents any of the entities in the list of Existing Nodes.
+
+        Existing Nodes:
+        {json.dumps(context['existing_nodes'], indent=2)}
+
+        New Node:
+        {json.dumps(context['extracted_nodes'], indent=2)}
+        Task:
+        1. If the New Node represents the same entity as any node in Existing Nodes, return 'is_duplicate: true' in the 
+            response. Otherwise, return 'is_duplicate: false'
+        2. If is_duplicate is true, also return the uuid of the existing node in the response
+        3. If is_duplicate is true, return a summary that synthesizes the information in the New Node summary and the 
+        summary of the Existing Node it is a duplicate of.
+
+        Guidelines:
+        1. Use both the name and summary of nodes to determine if the entities are duplicates, 
+            duplicate nodes may have different names
+
+        Respond with a JSON object in the following format:
+            {{
+                "is_duplicate": true or false,
+                "uuid": "uuid of the existing node like 5d643020624c42fa9de13f97b1b3fa39 or null",
+                "summary": "Brief summary of the node's role or significance. Takes information from the new and existing node"
+            }}
+        """,
+        ),
+    ]
+
+
 def node_list(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
@@ -128,18 +174,20 @@ def node_list(context: dict[str, Any]) -> list[Message]:
         {json.dumps(context['nodes'], indent=2)}
 
         Task:
-        1. Group nodes together such that all duplicate nodes are in the same list of names
-        2. All duplicate names should be grouped together in the same list
+        1. Group nodes together such that all duplicate nodes are in the same list of uuids
+        2. All duplicate uuids should be grouped together in the same list
+        3. Also return a new summary that synthesizes the summary into a new short summary
 
         Guidelines:
-        1. Each name from the list of nodes should appear EXACTLY once in your response
-        2. If a node has no duplicates, it should appear in the response in a list of only one name
+        1. Each uuid from the list of nodes should appear EXACTLY once in your response
+        2. If a node has no duplicates, it should appear in the response in a list of only one uuid
 
         Respond with a JSON object in the following format:
         {{
             "nodes": [
                 {{
-                    "names": ["myNode", "node that is a duplicate of myNode"],
+                    "uuids": ["5d643020624c42fa9de13f97b1b3fa39", "node that is a duplicate of 5d643020624c42fa9de13f97b1b3fa39"],
+                    "summary": "Brief summary of the node summaries that appear in the list of names."
                 }}
             ]
         }}
@@ -148,4 +196,4 @@ def node_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'v1': v1, 'v2': v2, 'node_list': node_list}
+versions: Versions = {'v1': v1, 'v2': v2, 'v3': v3, 'node_list': node_list}

graphiti_core/prompts/extract_edges.py

@@ -110,10 +110,11 @@ def v2(context: dict[str, Any]) -> list[Message]:
 
         Guidelines:
         1. Create edges only between the provided nodes.
-        2. Each edge should represent a clear relationship between two nodes.
+        2. Each edge should represent a clear relationship between two DISTINCT nodes.
         3. The relation_type should be a concise, all-caps description of the relationship (e.g., LOVES, IS_FRIENDS_WITH, WORKS_FOR).
         4. Provide a more detailed fact describing the relationship.
         5. Consider temporal aspects of relationships when relevant.
+        6. Avoid using the same node as the source and target of a relationship
 
         Respond with a JSON object in the following format:
         {{

graphiti_core/prompts/extract_nodes.py

@@ -55,6 +55,7 @@ def v1(context: dict[str, Any]) -> list[Message]:
         1. Focus on entities, concepts, or actors that are central to the current episode.
         2. Avoid creating nodes for relationships or actions (these will be handled as edges later).
         3. Provide a brief but informative summary for each node.
+        4. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
 
         Respond with a JSON object in the following format:
         {{
@@ -90,6 +91,7 @@ Guidelines:
 3. Provide concise but informative summaries for each extracted node.
 4. Avoid creating nodes for relationships or actions.
 5. Avoid creating nodes for temporal information like dates, times or years (these will be added to edges later).
+6. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
 
 Respond with a JSON object in the following format:
 {{

graphiti_core/search/search.py

@@ -63,12 +63,12 @@ class SearchResults(BaseModel):
 
 
 async def hybrid_search(
-        driver: AsyncDriver,
-        embedder,
-        query: str,
-        timestamp: datetime,
-        config: SearchConfig,
-        center_node_uuid: str | None = None,
+    driver: AsyncDriver,
+    embedder,
+    query: str,
+    timestamp: datetime,
+    config: SearchConfig,
+    center_node_uuid: str | None = None,
 ) -> SearchResults:
     start = time()
 
@@ -83,7 +83,7 @@ async def hybrid_search(
         nodes.extend(await get_mentioned_nodes(driver, episodes))
 
     if SearchMethod.bm25 in config.search_methods:
-        text_search = await edge_fulltext_search(query, driver, 2 * config.num_edges)
+        text_search = await edge_fulltext_search(driver, query, 2 * config.num_edges)
         search_results.append(text_search)
 
     if SearchMethod.cosine_similarity in config.search_methods:
@@ -95,7 +95,7 @@ async def hybrid_search(
         )
 
         similarity_search = await edge_similarity_search(
-            search_vector, driver, 2 * config.num_edges
+            driver, search_vector, 2 * config.num_edges
         )
         search_results.append(similarity_search)