graphiti-core 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

graphiti_core/edges.py

@@ -18,6 +18,7 @@ import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
 from time import time
+from typing import Any
 from uuid import uuid4
 
 from neo4j import AsyncDriver
@@ -32,6 +33,7 @@ logger = logging.getLogger(__name__)
 
 class Edge(BaseModel, ABC):
     uuid: str = Field(default_factory=lambda: uuid4().hex)
+    group_id: str | None = Field(description='partition of the graph')
     source_node_uuid: str
     target_node_uuid: str
     created_at: datetime
@@ -61,11 +63,12 @@ class EpisodicEdge(Edge):
         MATCH (episode:Episodic {uuid: $episode_uuid}) 
         MATCH (node:Entity {uuid: $entity_uuid}) 
         MERGE (episode)-[r:MENTIONS {uuid: $uuid}]->(node)
-        SET r = {uuid: $uuid, created_at: $created_at}
+        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
         RETURN r.uuid AS uuid""",
             episode_uuid=self.source_node_uuid,
             entity_uuid=self.target_node_uuid,
             uuid=self.uuid,
+            group_id=self.group_id,
             created_at=self.created_at,
         )
 
@@ -92,7 +95,8 @@ class EpisodicEdge(Edge):
             """
         MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
         RETURN
-            e.uuid As uuid, 
+            e.uuid As uuid,
+            e.group_id AS group_id,
             n.uuid AS source_node_uuid, 
             m.uuid AS target_node_uuid, 
             e.created_at AS created_at
@@ -100,17 +104,7 @@ class EpisodicEdge(Edge):
             uuid=uuid,
         )
 
-        edges: list[EpisodicEdge] = []
-
-        for record in records:
-            edges.append(
-                EpisodicEdge(
-                    uuid=record['uuid'],
-                    source_node_uuid=record['source_node_uuid'],
-                    target_node_uuid=record['target_node_uuid'],
-                    created_at=record['created_at'].to_native(),
-                )
-            )
+        edges = [get_episodic_edge_from_record(record) for record in records]
 
         logger.info(f'Found Edge: {uuid}')
 
@@ -153,7 +147,7 @@ class EntityEdge(Edge):
         MATCH (source:Entity {uuid: $source_uuid}) 
         MATCH (target:Entity {uuid: $target_uuid}) 
         MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target)
-        SET r = {uuid: $uuid, name: $name, fact: $fact, fact_embedding: $fact_embedding, 
+        SET r = {uuid: $uuid, name: $name, group_id: $group_id, fact: $fact, fact_embedding: $fact_embedding, 
         episodes: $episodes, created_at: $created_at, expired_at: $expired_at, 
         valid_at: $valid_at, invalid_at: $invalid_at}
         RETURN r.uuid AS uuid""",
@@ -161,6 +155,7 @@ class EntityEdge(Edge):
             target_uuid=self.target_node_uuid,
             uuid=self.uuid,
             name=self.name,
+            group_id=self.group_id,
             fact=self.fact,
             fact_embedding=self.fact_embedding,
             episodes=self.episodes,
@@ -198,6 +193,7 @@ class EntityEdge(Edge):
             m.uuid AS target_node_uuid,
             e.created_at AS created_at,
             e.name AS name,
+            e.group_id AS group_id,
             e.fact AS fact,
             e.fact_embedding AS fact_embedding,
             e.episodes AS episodes,
@@ -208,25 +204,36 @@ class EntityEdge(Edge):
             uuid=uuid,
         )
 
-        edges: list[EntityEdge] = []
-
-        for record in records:
-            edges.append(
-                EntityEdge(
-                    uuid=record['uuid'],
-                    source_node_uuid=record['source_node_uuid'],
-                    target_node_uuid=record['target_node_uuid'],
-                    fact=record['fact'],
-                    name=record['name'],
-                    episodes=record['episodes'],
-                    fact_embedding=record['fact_embedding'],
-                    created_at=record['created_at'].to_native(),
-                    expired_at=parse_db_date(record['expired_at']),
-                    valid_at=parse_db_date(record['valid_at']),
-                    invalid_at=parse_db_date(record['invalid_at']),
-                )
-            )
+        edges = [get_entity_edge_from_record(record) for record in records]
 
         logger.info(f'Found Edge: {uuid}')
 
         return edges[0]
+
+
+# Edge helpers
+def get_episodic_edge_from_record(record: Any) -> EpisodicEdge:
+    return EpisodicEdge(
+        uuid=record['uuid'],
+        group_id=record['group_id'],
+        source_node_uuid=record['source_node_uuid'],
+        target_node_uuid=record['target_node_uuid'],
+        created_at=record['created_at'].to_native(),
+    )
+
+
+def get_entity_edge_from_record(record: Any) -> EntityEdge:
+    return EntityEdge(
+        uuid=record['uuid'],
+        source_node_uuid=record['source_node_uuid'],
+        target_node_uuid=record['target_node_uuid'],
+        fact=record['fact'],
+        name=record['name'],
+        group_id=record['group_id'],
+        episodes=record['episodes'],
+        fact_embedding=record['fact_embedding'],
+        created_at=record['created_at'].to_native(),
+        expired_at=parse_db_date(record['expired_at']),
+        valid_at=parse_db_date(record['valid_at']),
+        invalid_at=parse_db_date(record['invalid_at']),
+    )

graphiti_core/graphiti.py

@@ -18,7 +18,6 @@ import asyncio
 import logging
 from datetime import datetime
 from time import time
-from typing import Callable
 
 from dotenv import load_dotenv
 from neo4j import AsyncGraphDatabase
@@ -59,11 +58,6 @@ 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,
-    prepare_edges_for_invalidation,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -125,7 +119,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -156,7 +150,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -183,6 +177,7 @@ class Graphiti:
         self,
         reference_time: datetime,
         last_n: int = EPISODE_WINDOW_LEN,
+        group_ids: list[str | None] | None = None,
     ) -> list[EpisodicNode]:
         """
         Retrieve the last n episodic nodes from the graph.
@@ -196,6 +191,8 @@ class Graphiti:
             The reference time to retrieve episodes before.
         last_n : int, optional
             The number of episodes to retrieve. Defaults to EPISODE_WINDOW_LEN.
+        group_ids : list[str | None], optional
+            The group ids to return data from.
 
         Returns
         -------
@@ -207,7 +204,7 @@ class Graphiti:
         The actual retrieval is performed by the `retrieve_episodes` function
         from the `graphiti_core.utils` module.
         """
-        return await retrieve_episodes(self.driver, reference_time, last_n)
+        return await retrieve_episodes(self.driver, reference_time, last_n, group_ids)
 
     async def add_episode(
         self,
@@ -216,8 +213,8 @@ class Graphiti:
         source_description: str,
         reference_time: datetime,
         source: EpisodeType = EpisodeType.message,
-        success_callback: Callable | None = None,
-        error_callback: Callable | None = None,
+        group_id: str | None = None,
+        uuid: str | None = None,
     ):
         """
         Process an episode and update the graph.
@@ -237,10 +234,10 @@ class Graphiti:
             The reference time for the episode.
         source : EpisodeType, optional
             The type of the episode. Defaults to EpisodeType.message.
-        success_callback : Callable | None, optional
-            A callback function to be called upon successful processing.
-        error_callback : Callable | None, optional
-            A callback function to be called if an error occurs during processing.
+        group_id : str | None
+            An id for the graph partition the episode is a part of.
+        uuid : str | None
+            Optional uuid of the episode.
 
         Returns
         -------
@@ -271,9 +268,12 @@ class Graphiti:
             embedder = self.llm_client.get_embedder()
             now = datetime.now()
 
-            previous_episodes = await self.retrieve_episodes(reference_time, last_n=3)
+            previous_episodes = await self.retrieve_episodes(
+                reference_time, last_n=3, group_ids=[group_id]
+            )
             episode = EpisodicNode(
                 name=name,
+                group_id=group_id,
                 labels=[],
                 source=source,
                 content=episode_body,
@@ -281,6 +281,7 @@ class Graphiti:
                 created_at=now,
                 valid_at=reference_time,
             )
+            episode.uuid = uuid if uuid is not None else episode.uuid
 
             # Extract entities as nodes
 
@@ -293,7 +294,7 @@ class Graphiti:
                 *[node.generate_name_embedding(embedder) for node in extracted_nodes]
             )
 
-            # Resolve extracted nodes with nodes already in the graph
+            # Resolve extracted nodes with nodes already in the graph and extract facts
             existing_nodes_lists: list[list[EntityNode]] = list(
                 await asyncio.gather(
                     *[get_relevant_nodes([node], self.driver) for node in extracted_nodes]
@@ -302,22 +303,29 @@ class Graphiti:
 
             logger.info(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
-            mentioned_nodes, _ = await resolve_extracted_nodes(
-                self.llm_client, extracted_nodes, existing_nodes_lists
+            (mentioned_nodes, uuid_map), extracted_edges = await asyncio.gather(
+                resolve_extracted_nodes(self.llm_client, extracted_nodes, existing_nodes_lists),
+                extract_edges(
+                    self.llm_client, episode, extracted_nodes, previous_episodes, group_id
+                ),
             )
             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, mentioned_nodes, previous_episodes
+            extracted_edges_with_resolved_pointers = resolve_edge_pointers(
+                extracted_edges, uuid_map
             )
 
             # calculate embeddings
-            await asyncio.gather(*[edge.generate_embedding(embedder) for edge in extracted_edges])
+            await asyncio.gather(
+                *[
+                    edge.generate_embedding(embedder)
+                    for edge in extracted_edges_with_resolved_pointers
+                ]
+            )
 
-            # Resolve extracted edges with edges already in the graph
-            existing_edges_list: list[list[EntityEdge]] = list(
+            # Resolve extracted edges with related edges already in the graph
+            related_edges_list: list[list[EntityEdge]] = list(
                 await asyncio.gather(
                     *[
                         get_relevant_edges(
@@ -327,80 +335,68 @@ class Graphiti:
                             edge.target_node_uuid,
                             RELEVANT_SCHEMA_LIMIT,
                         )
-                        for edge in extracted_edges
+                        for edge in extracted_edges_with_resolved_pointers
                     ]
                 )
             )
             logger.info(
-                f'Existing edges lists: {[(e.name, e.uuid) for edges_lst in existing_edges_list for e in edges_lst]}'
+                f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_list for e in edges_lst]}'
             )
-            logger.info(f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges]}')
-
-            deduped_edges: list[EntityEdge] = await resolve_extracted_edges(
-                self.llm_client, extracted_edges, existing_edges_list
+            logger.info(
+                f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges_with_resolved_pointers]}'
             )
 
-            # 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
-                ]
+            existing_source_edges_list: list[list[EntityEdge]] = list(
+                await asyncio.gather(
+                    *[
+                        get_relevant_edges(
+                            self.driver,
+                            [edge],
+                            edge.source_node_uuid,
+                            None,
+                            RELEVANT_SCHEMA_LIMIT,
+                        )
+                        for edge in extracted_edges_with_resolved_pointers
+                    ]
+                )
             )
 
-            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 is not None:
-                    edge.expired_at = now
-
-            entity_edges.extend(deduped_edges)
+            existing_target_edges_list: list[list[EntityEdge]] = list(
+                await asyncio.gather(
+                    *[
+                        get_relevant_edges(
+                            self.driver,
+                            [edge],
+                            None,
+                            edge.target_node_uuid,
+                            RELEVANT_SCHEMA_LIMIT,
+                        )
+                        for edge in extracted_edges_with_resolved_pointers
+                    ]
+                )
+            )
 
-            existing_edges: list[EntityEdge] = [
-                e for edge_lst in existing_edges_list for e in edge_lst
+            existing_edges_list: list[list[EntityEdge]] = [
+                source_lst + target_lst
+                for source_lst, target_lst in zip(
+                    existing_source_edges_list, existing_target_edges_list
+                )
             ]
 
-            (
-                old_edges_with_nodes_pending_invalidation,
-                new_edges_with_nodes,
-            ) = prepare_edges_for_invalidation(
-                existing_edges=existing_edges, new_edges=deduped_edges, nodes=nodes
-            )
-
-            invalidated_edges = await invalidate_edges(
+            resolved_edges, invalidated_edges = await resolve_extracted_edges(
                 self.llm_client,
-                old_edges_with_nodes_pending_invalidation,
-                new_edges_with_nodes,
+                extracted_edges_with_resolved_pointers,
+                related_edges_list,
+                existing_edges_list,
                 episode,
                 previous_episodes,
             )
 
-            for edge in invalidated_edges:
-                for existing_edge in existing_edges:
-                    if existing_edge.uuid == edge.uuid:
-                        existing_edge.expired_at = edge.expired_at
-                for deduped_edge in deduped_edges:
-                    if deduped_edge.uuid == edge.uuid:
-                        deduped_edge.expired_at = edge.expired_at
-            logger.info(f'Invalidated edges: {[(e.name, e.uuid) for e in invalidated_edges]}')
+            entity_edges.extend(resolved_edges + invalidated_edges)
 
-            entity_edges.extend(existing_edges)
+            logger.info(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_edges]}')
 
-            logger.info(f'Deduped edges: {[(e.name, e.uuid) for e in deduped_edges]}')
-
-            episodic_edges: list[EpisodicEdge] = build_episodic_edges(
-                mentioned_nodes,
-                episode,
-                now,
-            )
+            episodic_edges: list[EpisodicEdge] = build_episodic_edges(mentioned_nodes, episode, now)
 
             logger.info(f'Built episodic edges: {episodic_edges}')
 
@@ -413,18 +409,10 @@ class Graphiti:
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
 
-            if success_callback:
-                await success_callback(episode)
         except Exception as e:
-            if error_callback:
-                await error_callback(episode, e)
-            else:
-                raise e
+            raise e
 
-    async def add_episode_bulk(
-        self,
-        bulk_episodes: list[RawEpisode],
-    ):
+    async def add_episode_bulk(self, bulk_episodes: list[RawEpisode], group_id: str | None):
         """
         Process multiple episodes in bulk and update the graph.
 
@@ -435,6 +423,8 @@ class Graphiti:
         ----------
         bulk_episodes : list[RawEpisode]
             A list of RawEpisode objects to be processed and added to the graph.
+        group_id : str | None
+            An id for the graph partition the episode is a part of.
 
         Returns
         -------
@@ -471,6 +461,7 @@ class Graphiti:
                     source=episode.source,
                     content=episode.content,
                     source_description=episode.source_description,
+                    group_id=group_id,
                     created_at=now,
                     valid_at=episode.reference_time,
                 )
@@ -535,7 +526,13 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    async def search(self, query: str, center_node_uuid: str | None = None, num_results=10):
+    async def search(
+        self,
+        query: str,
+        center_node_uuid: str | None = None,
+        group_ids: list[str | None] | None = None,
+        num_results=10,
+    ):
         """
         Perform a hybrid search on the knowledge graph.
 
@@ -548,6 +545,8 @@ class Graphiti:
             The search query string.
         center_node_uuid: str, optional
             Facts will be reranked based on proximity to this node
+        group_ids : list[str | None] | None, optional
+            The graph partitions to return data from.
         num_results : int, optional
             The maximum number of results to return. Defaults to 10.
 
@@ -570,6 +569,7 @@ class Graphiti:
             num_episodes=0,
             num_edges=num_results,
             num_nodes=0,
+            group_ids=group_ids,
             search_methods=[SearchMethod.bm25, SearchMethod.cosine_similarity],
             reranker=reranker,
         )
@@ -598,7 +598,10 @@ class Graphiti:
         )
 
     async def get_nodes_by_query(
-        self, query: str, limit: int = RELEVANT_SCHEMA_LIMIT
+        self,
+        query: str,
+        group_ids: list[str | None] | None = None,
+        limit: int = RELEVANT_SCHEMA_LIMIT,
     ) -> list[EntityNode]:
         """
         Retrieve nodes from the graph database based on a text query.
@@ -610,6 +613,8 @@ class Graphiti:
         ----------
         query : str
             The text query to search for in the graph.
+        group_ids : list[str | None] | None, optional
+            The graph partitions to return data from.
         limit : int | None, optional
             The maximum number of results to return per search method.
             If None, a default limit will be applied.
@@ -634,5 +639,7 @@ class Graphiti:
         """
         embedder = self.llm_client.get_embedder()
         query_embedding = await generate_embedding(embedder, query)
-        relevant_nodes = await hybrid_node_search([query], [query_embedding], self.driver, limit)
+        relevant_nodes = await hybrid_node_search(
+            [query], [query_embedding], self.driver, group_ids, limit
+        )
         return relevant_nodes

graphiti_core/nodes.py

@@ -19,10 +19,10 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from enum import Enum
 from time import time
+from typing import Any
 from uuid import uuid4
 
 from neo4j import AsyncDriver
-from openai import OpenAI
 from pydantic import BaseModel, Field
 
 from graphiti_core.llm_client.config import EMBEDDING_DIM
@@ -69,6 +69,7 @@ class EpisodeType(Enum):
 class Node(BaseModel, ABC):
     uuid: str = Field(default_factory=lambda: uuid4().hex)
     name: str = Field(description='name of the node')
+    group_id: str | None = Field(description='partition of the graph')
     labels: list[str] = Field(default_factory=list)
     created_at: datetime = Field(default_factory=lambda: datetime.now())
 
@@ -106,11 +107,12 @@ class EpisodicNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Episodic {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, source_description: $source_description, source: $source, content: $content, 
+        SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content, 
         entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,
+            group_id=self.group_id,
             source_description=self.source_description,
             content=self.content,
             entity_edges=self.entity_edges,
@@ -141,29 +143,19 @@ class EpisodicNode(Node):
         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.source_description as source_description,
-            e.source as source
+            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
         """,
             uuid=uuid,
         )
 
-        episodes = [
-            EpisodicNode(
-                content=record['content'],
-                created_at=record['created_at'].to_native().timestamp(),
-                valid_at=(record['valid_at'].to_native()),
-                uuid=record['uuid'],
-                source=EpisodeType.from_str(record['source']),
-                name=record['name'],
-                source_description=record['source_description'],
-            )
-            for record in records
-        ]
+        episodes = [get_episodic_node_from_record(record) for record in records]
 
         logger.info(f'Found Node: {uuid}')
 
@@ -174,10 +166,6 @@ class EntityNode(Node):
     name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
     summary: str = Field(description='regional summary of surrounding edges', default_factory=str)
 
-    async def update_summary(self, driver: AsyncDriver): ...
-
-    async def refresh_summary(self, driver: AsyncDriver, llm_client: OpenAI): ...
-
     async def generate_name_embedding(self, embedder, model='text-embedding-3-small'):
         start = time()
         text = self.name.replace('\n', ' ')
@@ -192,10 +180,11 @@ class EntityNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Entity {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, summary: $summary, created_at: $created_at}
+        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, group_id: $group_id, summary: $summary, created_at: $created_at}
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,
+            group_id=self.group_id,
             summary=self.summary,
             name_embedding=self.name_embedding,
             created_at=self.created_at,
@@ -225,26 +214,43 @@ class EntityNode(Node):
         MATCH (n:Entity {uuid: $uuid})
         RETURN
             n.uuid As uuid, 
-            n.name AS name, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id
             n.created_at AS created_at, 
             n.summary AS summary
         """,
             uuid=uuid,
         )
 
-        nodes: list[EntityNode] = []
-
-        for record in records:
-            nodes.append(
-                EntityNode(
-                    uuid=record['uuid'],
-                    name=record['name'],
-                    labels=['Entity'],
-                    created_at=record['created_at'].to_native(),
-                    summary=record['summary'],
-                )
-            )
+        nodes = [get_entity_node_from_record(record) for record in records]
 
         logger.info(f'Found Node: {uuid}')
 
         return nodes[0]
+
+
+# Node helpers
+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()),
+        uuid=record['uuid'],
+        group_id=record['group_id'],
+        source=EpisodeType.from_str(record['source']),
+        name=record['name'],
+        source_description=record['source_description'],
+    )
+
+
+def get_entity_node_from_record(record: Any) -> EntityNode:
+    return EntityNode(
+        uuid=record['uuid'],
+        name=record['name'],
+        group_id=record['group_id'],
+        name_embedding=record['name_embedding'],
+        labels=['Entity'],
+        created_at=record['created_at'].to_native(),
+        summary=record['summary'],
+    )

graphiti_core/prompts/dedupe_edges.py

@@ -129,7 +129,7 @@ def v3(context: dict[str, Any]) -> list[Message]:
         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)}
+        {json.dumps(context['related_edges'], indent=2)}
 
         New Edge:
         {json.dumps(context['extracted_edges'], indent=2)}