graphiti-core 0.2.2__tar.gz → 0.2.3__tar.gz

{graphiti_core-0.2.2 → graphiti_core-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.2.2
+Version: 0.2.3
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -12,11 +12,10 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: diskcache (>=5.6.3,<6.0.0)
-Requires-Dist: fastapi (>=0.112.0,<0.113.0)
 Requires-Dist: neo4j (>=5.23.0,<6.0.0)
+Requires-Dist: numpy (>=2.1.1,<3.0.0)
 Requires-Dist: openai (>=1.38.0,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
-Requires-Dist: sentence-transformers (>=3.0.1,<4.0.0)
 Requires-Dist: tenacity (<9.0.0)
 Description-Content-Type: text/markdown
 

{graphiti_core-0.2.2 → graphiti_core-0.2.3}/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-0.2.2 → graphiti_core-0.2.3}/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
@@ -120,7 +119,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -151,7 +150,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -178,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.
@@ -191,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
         -------
@@ -202,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,
@@ -211,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.
@@ -232,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
         -------
@@ -266,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,
@@ -276,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
 
@@ -299,7 +305,9 @@ class Graphiti:
 
             (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),
+                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)
@@ -388,11 +396,7 @@ class Graphiti:
 
             logger.info(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_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}')
 
@@ -405,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.
 
@@ -427,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
         -------
@@ -463,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,
                 )
@@ -527,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.
 
@@ -540,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.
 
@@ -562,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,
         )
@@ -590,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.
@@ -602,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.
@@ -626,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-0.2.2 → graphiti_core-0.2.3}/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,
@@ -227,25 +216,14 @@ class EntityNode(Node):
             n.uuid As uuid, 
             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'],
-                    name_embedding=record['name_embedding'],
-                    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}')
 
@@ -253,3 +231,26 @@ class EntityNode(Node):
 
 
 # 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-0.2.2 → graphiti_core-0.2.3}/graphiti_core/search/search.py

@@ -52,6 +52,7 @@ class SearchConfig(BaseModel):
     num_edges: int = Field(default=10)
     num_nodes: int = Field(default=10)
     num_episodes: int = EPISODE_WINDOW_LEN
+    group_ids: list[str | None] | None
     search_methods: list[SearchMethod]
     reranker: Reranker | None
 
@@ -83,7 +84,9 @@ 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(driver, query, None, None, 2 * config.num_edges)
+        text_search = await edge_fulltext_search(
+            driver, query, None, None, config.group_ids, 2 * config.num_edges
+        )
         search_results.append(text_search)
 
     if SearchMethod.cosine_similarity in config.search_methods:
@@ -95,7 +98,7 @@ async def hybrid_search(
         )
 
         similarity_search = await edge_similarity_search(
-            driver, search_vector, None, None, 2 * config.num_edges
+            driver, search_vector, None, None, config.group_ids, 2 * config.num_edges
         )
         search_results.append(similarity_search)