graphiti-core 0.3.0__tar.gz → 0.3.1__tar.gz

{graphiti_core-0.3.0 → graphiti_core-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.0
+Version: 0.3.1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

graphiti_core-0.3.1/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-0.3.0 → graphiti_core-0.3.1}/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,
@@ -548,7 +553,7 @@ class Graphiti:
         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.
@@ -564,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
@@ -581,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,
             )
@@ -606,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.
@@ -629,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
@@ -655,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-0.3.1/graphiti_core/helpers.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.
+"""
+
+from datetime import datetime
+
+from neo4j import time as neo4j_time
+
+
+def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
+    return neo_date.to_native() if neo_date else None

graphiti_core-0.3.1/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-0.3.1/graphiti_core/llm_client/utils.py

@@ -0,0 +1,38 @@
+"""
+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
+
+from graphiti_core.llm_client.config import EMBEDDING_DIM
+
+logger = logging.getLogger(__name__)
+
+
+async def generate_embedding(
+    embedder: typing.Any, text: str, model: str = 'text-embedding-3-small'
+):
+    start = time()
+
+    text = text.replace('\n', ' ')
+    embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
+    embedding = embedding[:EMBEDDING_DIM]
+
+    end = time()
+    logger.debug(f'embedded text of length {len(text)} in {end - start} ms')
+
+    return embedding

{graphiti_core-0.3.0 → graphiti_core-0.3.1}/graphiti_core/prompts/extract_edge_dates.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 typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion

graphiti_core-0.3.1/graphiti_core/search/search.py

@@ -0,0 +1,242 @@
+"""
+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
+from time import time
+
+from neo4j import AsyncDriver
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.errors import SearchRerankerError
+from graphiti_core.llm_client.config import EMBEDDING_DIM
+from graphiti_core.nodes import CommunityNode, EntityNode
+from graphiti_core.search.search_config import (
+    DEFAULT_SEARCH_LIMIT,
+    CommunityReranker,
+    CommunitySearchConfig,
+    CommunitySearchMethod,
+    EdgeReranker,
+    EdgeSearchConfig,
+    EdgeSearchMethod,
+    NodeReranker,
+    NodeSearchConfig,
+    NodeSearchMethod,
+    SearchConfig,
+    SearchResults,
+)
+from graphiti_core.search.search_utils import (
+    community_fulltext_search,
+    community_similarity_search,
+    edge_fulltext_search,
+    edge_similarity_search,
+    node_distance_reranker,
+    node_fulltext_search,
+    node_similarity_search,
+    rrf,
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def search(
+    driver: AsyncDriver,
+    embedder,
+    query: str,
+    group_ids: list[str | None] | None,
+    config: SearchConfig,
+    center_node_uuid: str | None = None,
+) -> SearchResults:
+    start = time()
+    query = query.replace('\n', ' ')
+    # if group_ids is empty, set it to None
+    group_ids = group_ids if group_ids else None
+    edges = (
+        await edge_search(
+            driver, embedder, query, group_ids, config.edge_config, center_node_uuid, config.limit
+        )
+        if config.edge_config is not None
+        else []
+    )
+    nodes = (
+        await node_search(
+            driver, embedder, query, group_ids, config.node_config, center_node_uuid, config.limit
+        )
+        if config.node_config is not None
+        else []
+    )
+    communities = (
+        await community_search(
+            driver, embedder, query, group_ids, config.community_config, config.limit
+        )
+        if config.community_config is not None
+        else []
+    )
+
+    results = SearchResults(
+        edges=edges[: config.limit],
+        nodes=nodes[: config.limit],
+        communities=communities[: config.limit],
+    )
+
+    end = time()
+
+    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
+
+    return results
+
+
+async def edge_search(
+    driver: AsyncDriver,
+    embedder,
+    query: str,
+    group_ids: list[str | None] | None,
+    config: EdgeSearchConfig,
+    center_node_uuid: str | None = None,
+    limit=DEFAULT_SEARCH_LIMIT,
+) -> list[EntityEdge]:
+    search_results: list[list[EntityEdge]] = []
+
+    if EdgeSearchMethod.bm25 in config.search_methods:
+        text_search = await edge_fulltext_search(driver, query, None, None, group_ids, 2 * limit)
+        search_results.append(text_search)
+
+    if EdgeSearchMethod.cosine_similarity in config.search_methods:
+        search_vector = (
+            (await embedder.create(input=[query], model='text-embedding-3-small'))
+            .data[0]
+            .embedding[:EMBEDDING_DIM]
+        )
+
+        similarity_search = await edge_similarity_search(
+            driver, search_vector, None, None, group_ids, 2 * limit
+        )
+        search_results.append(similarity_search)
+
+    if len(search_results) > 1 and config.reranker is None:
+        raise SearchRerankerError('Multiple edge searches enabled without a reranker')
+
+    edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
+
+    reranked_uuids: list[str] = []
+    if config.reranker == EdgeReranker.rrf:
+        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
+
+        reranked_uuids = rrf(search_result_uuids)
+    elif config.reranker == EdgeReranker.node_distance:
+        if center_node_uuid is None:
+            raise SearchRerankerError('No center node provided for Node Distance reranker')
+
+        source_to_edge_uuid_map = {
+            edge.source_node_uuid: edge.uuid for result in search_results for edge in result
+        }
+        source_uuids = [[edge.source_node_uuid for edge in result] for result in search_results]
+
+        reranked_node_uuids = await node_distance_reranker(driver, source_uuids, center_node_uuid)
+
+        reranked_uuids = [source_to_edge_uuid_map[node_uuid] for node_uuid in reranked_node_uuids]
+
+    reranked_edges = [edge_uuid_map[uuid] for uuid in reranked_uuids]
+
+    return reranked_edges
+
+
+async def node_search(
+    driver: AsyncDriver,
+    embedder,
+    query: str,
+    group_ids: list[str | None] | None,
+    config: NodeSearchConfig,
+    center_node_uuid: str | None = None,
+    limit=DEFAULT_SEARCH_LIMIT,
+) -> list[EntityNode]:
+    search_results: list[list[EntityNode]] = []
+
+    if NodeSearchMethod.bm25 in config.search_methods:
+        text_search = await node_fulltext_search(driver, query, group_ids, 2 * limit)
+        search_results.append(text_search)
+
+    if NodeSearchMethod.cosine_similarity in config.search_methods:
+        search_vector = (
+            (await embedder.create(input=[query], model='text-embedding-3-small'))
+            .data[0]
+            .embedding[:EMBEDDING_DIM]
+        )
+
+        similarity_search = await node_similarity_search(
+            driver, search_vector, group_ids, 2 * limit
+        )
+        search_results.append(similarity_search)
+
+    if len(search_results) > 1 and config.reranker is None:
+        raise SearchRerankerError('Multiple node searches enabled without a reranker')
+
+    search_result_uuids = [[node.uuid for node in result] for result in search_results]
+    node_uuid_map = {node.uuid: node for result in search_results for node in result}
+
+    reranked_uuids: list[str] = []
+    if config.reranker == NodeReranker.rrf:
+        reranked_uuids = rrf(search_result_uuids)
+    elif config.reranker == NodeReranker.node_distance:
+        if center_node_uuid is None:
+            raise SearchRerankerError('No center node provided for Node Distance reranker')
+        reranked_uuids = await node_distance_reranker(driver, search_result_uuids, center_node_uuid)
+
+    reranked_nodes = [node_uuid_map[uuid] for uuid in reranked_uuids]
+
+    return reranked_nodes
+
+
+async def community_search(
+    driver: AsyncDriver,
+    embedder,
+    query: str,
+    group_ids: list[str | None] | None,
+    config: CommunitySearchConfig,
+    limit=DEFAULT_SEARCH_LIMIT,
+) -> list[CommunityNode]:
+    search_results: list[list[CommunityNode]] = []
+
+    if CommunitySearchMethod.bm25 in config.search_methods:
+        text_search = await community_fulltext_search(driver, query, group_ids, 2 * limit)
+        search_results.append(text_search)
+
+    if CommunitySearchMethod.cosine_similarity in config.search_methods:
+        search_vector = (
+            (await embedder.create(input=[query], model='text-embedding-3-small'))
+            .data[0]
+            .embedding[:EMBEDDING_DIM]
+        )
+
+        similarity_search = await community_similarity_search(
+            driver, search_vector, group_ids, 2 * limit
+        )
+        search_results.append(similarity_search)
+
+    if len(search_results) > 1 and config.reranker is None:
+        raise SearchRerankerError('Multiple node searches enabled without a reranker')
+
+    search_result_uuids = [[community.uuid for community in result] for result in search_results]
+    community_uuid_map = {
+        community.uuid: community for result in search_results for community in result
+    }
+
+    reranked_uuids: list[str] = []
+    if config.reranker == CommunityReranker.rrf:
+        reranked_uuids = rrf(search_result_uuids)
+
+    reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
+
+    return reranked_communities

graphiti_core-0.3.1/graphiti_core/search/search_config.py

@@ -0,0 +1,81 @@
+"""
+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 enum import Enum
+
+from pydantic import BaseModel, Field
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.nodes import CommunityNode, EntityNode
+
+DEFAULT_SEARCH_LIMIT = 10
+
+
+class EdgeSearchMethod(Enum):
+    cosine_similarity = 'cosine_similarity'
+    bm25 = 'bm25'
+
+
+class NodeSearchMethod(Enum):
+    cosine_similarity = 'cosine_similarity'
+    bm25 = 'bm25'
+
+
+class CommunitySearchMethod(Enum):
+    cosine_similarity = 'cosine_similarity'
+    bm25 = 'bm25'
+
+
+class EdgeReranker(Enum):
+    rrf = 'reciprocal_rank_fusion'
+    node_distance = 'node_distance'
+
+
+class NodeReranker(Enum):
+    rrf = 'reciprocal_rank_fusion'
+    node_distance = 'node_distance'
+
+
+class CommunityReranker(Enum):
+    rrf = 'reciprocal_rank_fusion'
+
+
+class EdgeSearchConfig(BaseModel):
+    search_methods: list[EdgeSearchMethod]
+    reranker: EdgeReranker | None
+
+
+class NodeSearchConfig(BaseModel):
+    search_methods: list[NodeSearchMethod]
+    reranker: NodeReranker | None
+
+
+class CommunitySearchConfig(BaseModel):
+    search_methods: list[CommunitySearchMethod]
+    reranker: CommunityReranker | None
+
+
+class SearchConfig(BaseModel):
+    edge_config: EdgeSearchConfig | None = Field(default=None)
+    node_config: NodeSearchConfig | None = Field(default=None)
+    community_config: CommunitySearchConfig | None = Field(default=None)
+    limit: int = Field(default=DEFAULT_SEARCH_LIMIT)
+
+
+class SearchResults(BaseModel):
+    edges: list[EntityEdge]
+    nodes: list[EntityNode]
+    communities: list[CommunityNode]

graphiti_core-0.3.1/graphiti_core/search/search_config_recipes.py

@@ -0,0 +1,84 @@
+"""
+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 graphiti_core.search.search_config import (
+    CommunityReranker,
+    CommunitySearchConfig,
+    CommunitySearchMethod,
+    EdgeReranker,
+    EdgeSearchConfig,
+    EdgeSearchMethod,
+    NodeReranker,
+    NodeSearchConfig,
+    NodeSearchMethod,
+    SearchConfig,
+)
+
+# Performs a hybrid search with rrf reranking over edges, nodes, and communities
+COMBINED_HYBRID_SEARCH_RRF = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
+        reranker=EdgeReranker.rrf,
+    ),
+    node_config=NodeSearchConfig(
+        search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
+        reranker=NodeReranker.rrf,
+    ),
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.rrf,
+    ),
+)
+
+# performs a hybrid search over edges with rrf reranking
+EDGE_HYBRID_SEARCH_RRF = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
+        reranker=EdgeReranker.rrf,
+    )
+)
+
+# performs a hybrid search over edges with node distance reranking
+EDGE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
+        reranker=EdgeReranker.node_distance,
+    )
+)
+
+# performs a hybrid search over nodes with rrf reranking
+NODE_HYBRID_SEARCH_RRF = SearchConfig(
+    node_config=NodeSearchConfig(
+        search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
+        reranker=NodeReranker.rrf,
+    )
+)
+
+# performs a hybrid search over nodes with node distance reranking
+NODE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
+    node_config=NodeSearchConfig(
+        search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
+        reranker=NodeReranker.node_distance,
+    )
+)
+
+# performs a hybrid search over communities with rrf reranking
+COMMUNITY_HYBRID_SEARCH_RRF = SearchConfig(
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.rrf,
+    )
+)