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

graphiti_core/edges.py

@@ -24,6 +24,7 @@ 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
@@ -41,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)
@@ -76,19 +87,6 @@ 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(
@@ -107,7 +105,8 @@ class EpisodicEdge(Edge):
         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]
 
 
@@ -169,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(
@@ -206,6 +192,49 @@ class EntityEdge(Edge):
 
         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]
@@ -237,3 +266,13 @@ def get_entity_edge_from_record(record: Any) -> EntityEdge:
         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,43 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+
+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)
+
+
+class SearchRerankerError(GraphitiError):
+    """Raised when a node is not found."""
+
+    def __init__(self, text: str):
+        self.message = text
+        super().__init__(self.message)

graphiti_core/graphiti.py

@@ -24,14 +24,19 @@ from neo4j import AsyncGraphDatabase
 
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.llm_client import LLMClient, OpenAIClient
-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 import SearchConfig, search
+from graphiti_core.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults
+from graphiti_core.search.search_config_recipes import (
+    EDGE_HYBRID_SEARCH_NODE_DISTANCE,
+    EDGE_HYBRID_SEARCH_RRF,
+    NODE_HYBRID_SEARCH_NODE_DISTANCE,
+    NODE_HYBRID_SEARCH_RRF,
+)
 from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
     get_relevant_edges,
     get_relevant_nodes,
-    hybrid_node_search,
 )
 from graphiti_core.utils import (
     build_episodic_edges,
@@ -46,6 +51,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,
@@ -412,7 +421,7 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    async def add_episode_bulk(self, bulk_episodes: list[RawEpisode], group_id: str | None):
+    async def add_episode_bulk(self, bulk_episodes: list[RawEpisode], group_id: str | None = None):
         """
         Process multiple episodes in bulk and update the graph.
 
@@ -526,12 +535,25 @@ class Graphiti:
         except Exception as e:
             raise e
 
+    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,
+        num_results=DEFAULT_SEARCH_LIMIT,
     ):
         """
         Perform a hybrid search on the knowledge graph.
@@ -547,7 +569,7 @@ class Graphiti:
             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
+        limit : int, optional
             The maximum number of results to return. Defaults to 10.
 
         Returns
@@ -564,21 +586,17 @@ class Graphiti:
         The search is performed using the current date and time as the reference
         point for temporal relevance.
         """
-        reranker = Reranker.rrf if center_node_uuid is None else Reranker.node_distance
-        search_config = SearchConfig(
-            num_episodes=0,
-            num_edges=num_results,
-            num_nodes=0,
-            group_ids=group_ids,
-            search_methods=[SearchMethod.bm25, SearchMethod.cosine_similarity],
-            reranker=reranker,
+        search_config = (
+            EDGE_HYBRID_SEARCH_RRF if center_node_uuid is None else EDGE_HYBRID_SEARCH_NODE_DISTANCE
         )
+        search_config.limit = num_results
+
         edges = (
-            await hybrid_search(
+            await search(
                 self.driver,
                 self.llm_client.get_embedder(),
                 query,
-                datetime.now(),
+                group_ids,
                 search_config,
                 center_node_uuid,
             )
@@ -589,19 +607,20 @@ class Graphiti:
     async def _search(
         self,
         query: str,
-        timestamp: datetime,
         config: SearchConfig,
+        group_ids: list[str | None] | None = None,
         center_node_uuid: str | None = None,
-    ):
-        return await hybrid_search(
-            self.driver, self.llm_client.get_embedder(), query, timestamp, config, center_node_uuid
+    ) -> SearchResults:
+        return await search(
+            self.driver, self.llm_client.get_embedder(), query, group_ids, config, center_node_uuid
         )
 
     async def get_nodes_by_query(
         self,
         query: str,
+        center_node_uuid: str | None = None,
         group_ids: list[str | None] | None = None,
-        limit: int = RELEVANT_SCHEMA_LIMIT,
+        limit: int = DEFAULT_SEARCH_LIMIT,
     ) -> list[EntityNode]:
         """
         Retrieve nodes from the graph database based on a text query.
@@ -612,7 +631,9 @@ class Graphiti:
         Parameters
         ----------
         query : str
-            The text query to search for in the graph.
+            The text query to search for in the graph
+        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.
         limit : int | None, optional
@@ -638,8 +659,12 @@ class Graphiti:
         If not specified, a default limit (defined in the search functions) will be used.
         """
         embedder = self.llm_client.get_embedder()
-        query_embedding = await generate_embedding(embedder, query)
-        relevant_nodes = await hybrid_node_search(
-            [query], [query_embedding], self.driver, group_ids, limit
+        search_config = (
+            NODE_HYBRID_SEARCH_RRF if center_node_uuid is None else NODE_HYBRID_SEARCH_NODE_DISTANCE
         )
-        return relevant_nodes
+        search_config.limit = limit
+
+        nodes = (
+            await search(self.driver, embedder, query, group_ids, search_config, center_node_uuid)
+        ).nodes
+        return nodes

graphiti_core/helpers.py

@@ -1,3 +1,19 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 from datetime import datetime
 
 from neo4j import time as neo4j_time

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,23 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+
+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

graphiti_core/llm_client/utils.py

@@ -1,3 +1,19 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 import logging
 import typing
 from time import time
@@ -17,6 +33,6 @@ async def generate_embedding(
     embedding = embedding[:EMBEDDING_DIM]
 
     end = time()
-    logger.debug(f'embedded text of length {len(text)} in {end-start} ms')
+    logger.debug(f'embedded text of length {len(text)} in {end - start} ms')
 
     return embedding