graphiti-core 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

graphiti_core/edges.py

@@ -18,11 +18,13 @@ 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
 from pydantic import BaseModel, Field
 
+from graphiti_core.errors import EdgeNotFoundError
 from graphiti_core.helpers import parse_db_date
 from graphiti_core.llm_client.config import EMBEDDING_DIM
 from graphiti_core.nodes import Node
@@ -32,6 +34,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
@@ -39,8 +42,18 @@ class Edge(BaseModel, ABC):
     @abstractmethod
     async def save(self, driver: AsyncDriver): ...
 
-    @abstractmethod
-    async def delete(self, driver: AsyncDriver): ...
+    async def delete(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MATCH (n)-[e {uuid: $uuid}]->(m)
+        DELETE e
+        """,
+            uuid=self.uuid,
+        )
+
+        logger.info(f'Deleted Edge: {self.uuid}')
+
+        return result
 
     def __hash__(self):
         return hash(self.uuid)
@@ -61,11 +74,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,
         )
 
@@ -73,26 +87,14 @@ class EpisodicEdge(Edge):
 
         return result
 
-    async def delete(self, driver: AsyncDriver):
-        result = await driver.execute_query(
-            """
-        MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
-        DELETE e
-        """,
-            uuid=self.uuid,
-        )
-
-        logger.info(f'Deleted Edge: {self.uuid}')
-
-        return result
-
     @classmethod
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         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,20 +102,11 @@ 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}')
-
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuid)
         return edges[0]
 
 
@@ -153,7 +146,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 +154,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,
@@ -174,19 +168,6 @@ class EntityEdge(Edge):
 
         return result
 
-    async def delete(self, driver: AsyncDriver):
-        result = await driver.execute_query(
-            """
-        MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
-        DELETE e
-        """,
-            uuid=self.uuid,
-        )
-
-        logger.info(f'Deleted Edge: {self.uuid}')
-
-        return result
-
     @classmethod
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
         records, _, _ = await driver.execute_query(
@@ -198,6 +179,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 +190,89 @@ 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}')
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuid)
+        return edges[0]
+
+
+class CommunityEdge(Edge):
+    async def save(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MATCH (community:Community {uuid: $community_uuid}) 
+        MATCH (node:Entity | Community {uuid: $entity_uuid}) 
+        MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node)
+        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+        RETURN r.uuid AS uuid""",
+            community_uuid=self.source_node_uuid,
+            entity_uuid=self.target_node_uuid,
+            uuid=self.uuid,
+            group_id=self.group_id,
+            created_at=self.created_at,
+        )
+
+        logger.info(f'Saved edge to neo4j: {self.uuid}')
+
+        return result
+
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m:Entity | Community)
+        RETURN
+            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
+        """,
+            uuid=uuid,
+        )
+
+        edges = [get_community_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']),
+    )
+
+
+def get_community_edge_from_record(record: Any):
+    return CommunityEdge(
+        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(),
+    )

graphiti_core/errors.py

@@ -0,0 +1,18 @@
+class GraphitiError(Exception):
+    """Base exception class for Graphiti Core."""
+
+
+class EdgeNotFoundError(GraphitiError):
+    """Raised when an edge is not found."""
+
+    def __init__(self, uuid: str):
+        self.message = f'edge {uuid} not found'
+        super().__init__(self.message)
+
+
+class NodeNotFoundError(GraphitiError):
+    """Raised when a node is not found."""
+
+    def __init__(self, uuid: str):
+        self.message = f'node {uuid} not found'
+        super().__init__(self.message)

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
@@ -47,6 +46,10 @@ from graphiti_core.utils.bulk_utils import (
     resolve_edge_pointers,
     retrieve_previous_episodes_bulk,
 )
+from graphiti_core.utils.maintenance.community_operations import (
+    build_communities,
+    remove_communities,
+)
 from graphiti_core.utils.maintenance.edge_operations import (
     extract_edges,
     resolve_extracted_edges,
@@ -120,7 +123,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -151,7 +154,7 @@ class Graphiti:
 
         Parameters
         ----------
-        None
+        self
 
         Returns
         -------
@@ -178,6 +181,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 +195,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 +208,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 +217,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 +238,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 +272,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 +285,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 +309,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 +400,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 +413,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 = None):
         """
         Process multiple episodes in bulk and update the graph.
 
@@ -427,6 +427,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 +465,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 +530,26 @@ 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 build_communities(self):
+        embedder = self.llm_client.get_embedder()
+
+        # Clear existing communities
+        await remove_communities(self.driver)
+
+        community_nodes, community_edges = await build_communities(self.driver, self.llm_client)
+
+        await asyncio.gather(*[node.generate_name_embedding(embedder) for node in community_nodes])
+
+        await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
+        await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
+
+    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 +562,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 +586,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 +615,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 +630,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 +656,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/llm_client/__init__.py

@@ -1,5 +1,6 @@
 from .client import LLMClient
 from .config import LLMConfig
+from .errors import RateLimitError
 from .openai_client import OpenAIClient
 
-__all__ = ['LLMClient', 'OpenAIClient', 'LLMConfig']
+__all__ = ['LLMClient', 'OpenAIClient', 'LLMConfig', 'RateLimitError']

graphiti_core/llm_client/anthropic_client.py

@@ -18,12 +18,14 @@ import json
 import logging
 import typing
 
+import anthropic
 from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
 
 from ..prompts.models import Message
 from .client import LLMClient
 from .config import LLMConfig
+from .errors import RateLimitError
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +37,11 @@ class AnthropicClient(LLMClient):
         if config is None:
             config = LLMConfig()
         super().__init__(config, cache)
-        self.client = AsyncAnthropic(api_key=config.api_key)
+        self.client = AsyncAnthropic(
+            api_key=config.api_key,
+            # we'll use tenacity to retry
+            max_retries=1,
+        )
 
     def get_embedder(self) -> typing.Any:
         openai_client = AsyncOpenAI()
@@ -58,6 +64,8 @@ class AnthropicClient(LLMClient):
             )
 
             return json.loads('{' + result.content[0].text)  # type: ignore
+        except anthropic.RateLimitError as e:
+            raise RateLimitError from e
         except Exception as e:
             logger.error(f'Error in generating LLM response: {e}')
             raise

graphiti_core/llm_client/client.py

@@ -22,10 +22,11 @@ from abc import ABC, abstractmethod
 
 import httpx
 from diskcache import Cache
-from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
 
 from ..prompts.models import Message
 from .config import LLMConfig
+from .errors import RateLimitError
 
 DEFAULT_TEMPERATURE = 0
 DEFAULT_CACHE_DIR = './llm_cache'
@@ -33,7 +34,10 @@ DEFAULT_CACHE_DIR = './llm_cache'
 logger = logging.getLogger(__name__)
 
 
-def is_server_error(exception):
+def is_server_or_retry_error(exception):
+    if isinstance(exception, RateLimitError):
+        return True
+
     return (
         isinstance(exception, httpx.HTTPStatusError) and 500 <= exception.response.status_code < 600
     )
@@ -56,18 +60,21 @@ class LLMClient(ABC):
         pass
 
     @retry(
-        stop=stop_after_attempt(3),
-        wait=wait_exponential(multiplier=1, min=4, max=10),
-        retry=retry_if_exception(is_server_error),
+        stop=stop_after_attempt(4),
+        wait=wait_random_exponential(multiplier=10, min=5, max=120),
+        retry=retry_if_exception(is_server_or_retry_error),
+        after=lambda retry_state: logger.warning(
+            f'Retrying {retry_state.fn.__name__ if retry_state.fn else "function"} after {retry_state.attempt_number} attempts...'
+        )
+        if retry_state.attempt_number > 1
+        else None,
+        reraise=True,
     )
     async def _generate_response_with_retry(self, messages: list[Message]) -> dict[str, typing.Any]:
         try:
             return await self._generate_response(messages)
-        except httpx.HTTPStatusError as e:
-            if not is_server_error(e):
-                raise Exception(f'LLM request error: {e}') from e
-            else:
-                raise
+        except (httpx.HTTPStatusError, RateLimitError) as e:
+            raise e
 
     @abstractmethod
     async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:

graphiti_core/llm_client/errors.py

@@ -0,0 +1,6 @@
+class RateLimitError(Exception):
+    """Exception raised when the rate limit is exceeded."""
+
+    def __init__(self, message='Rate limit exceeded. Please try again later.'):
+        self.message = message
+        super().__init__(self.message)

graphiti_core/llm_client/groq_client.py

@@ -18,6 +18,7 @@ import json
 import logging
 import typing
 
+import groq
 from groq import AsyncGroq
 from groq.types.chat import ChatCompletionMessageParam
 from openai import AsyncOpenAI
@@ -25,6 +26,7 @@ from openai import AsyncOpenAI
 from ..prompts.models import Message
 from .client import LLMClient
 from .config import LLMConfig
+from .errors import RateLimitError
 
 logger = logging.getLogger(__name__)
 
@@ -59,6 +61,8 @@ class GroqClient(LLMClient):
             )
             result = response.choices[0].message.content or ''
             return json.loads(result)
+        except groq.RateLimitError as e:
+            raise RateLimitError from e
         except Exception as e:
             logger.error(f'Error in generating LLM response: {e}')
             raise

graphiti_core/llm_client/openai_client.py

@@ -18,12 +18,14 @@ import json
 import logging
 import typing
 
+import openai
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletionMessageParam
 
 from ..prompts.models import Message
 from .client import LLMClient
 from .config import LLMConfig
+from .errors import RateLimitError
 
 logger = logging.getLogger(__name__)
 
@@ -59,6 +61,8 @@ class OpenAIClient(LLMClient):
             )
             result = response.choices[0].message.content or ''
             return json.loads(result)
+        except openai.RateLimitError as e:
+            raise RateLimitError from e
         except Exception as e:
             logger.error(f'Error in generating LLM response: {e}')
             raise