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

graphiti_core/search/search.py

@@ -15,131 +15,228 @@ limitations under the License.
 """
 
 import logging
-from datetime import datetime
-from enum import Enum
 from time import time
 
 from neo4j import AsyncDriver
-from pydantic import BaseModel, Field
 
 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 EntityNode, EpisodicNode
+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,
-    get_mentioned_nodes,
     node_distance_reranker,
+    node_fulltext_search,
+    node_similarity_search,
     rrf,
 )
-from graphiti_core.utils import retrieve_episodes
-from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
 
 logger = logging.getLogger(__name__)
 
 
-class SearchMethod(Enum):
-    cosine_similarity = 'cosine_similarity'
-    bm25 = 'bm25'
-
-
-class Reranker(Enum):
-    rrf = 'reciprocal_rank_fusion'
-    node_distance = 'node_distance'
+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],
+    )
 
-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
+    end = time()
 
+    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
 
-class SearchResults(BaseModel):
-    episodes: list[EpisodicNode]
-    nodes: list[EntityNode]
-    edges: list[EntityEdge]
+    return results
 
 
-async def hybrid_search(
+async def edge_search(
     driver: AsyncDriver,
     embedder,
     query: str,
-    timestamp: datetime,
-    config: SearchConfig,
+    group_ids: list[str | None] | None,
+    config: EdgeSearchConfig,
     center_node_uuid: str | None = None,
-) -> SearchResults:
-    start = time()
-
-    episodes = []
-    nodes = []
-    edges = []
+    limit=DEFAULT_SEARCH_LIMIT,
+) -> list[EntityEdge]:
+    search_results: list[list[EntityEdge]] = []
 
-    search_results = []
+    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 config.num_episodes > 0:
-        episodes.extend(await retrieve_episodes(driver, timestamp, config.num_episodes))
-        nodes.extend(await get_mentioned_nodes(driver, episodes))
+    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]
+        )
 
-    if SearchMethod.bm25 in config.search_methods:
-        text_search = await edge_fulltext_search(
-            driver, query, None, None, config.group_ids, 2 * config.num_edges
+        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 SearchMethod.cosine_similarity in config.search_methods:
-        query_text = query.replace('\n', ' ')
+    if NodeSearchMethod.cosine_similarity in config.search_methods:
         search_vector = (
-            (await embedder.create(input=[query_text], model='text-embedding-3-small'))
+            (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, config.group_ids, 2 * config.num_edges
+        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:
-        logger.exception('Multiple searches enabled without a reranker')
-        raise Exception('Multiple searches enabled without a reranker')
+        raise SearchRerankerError('Multiple node searches enabled without a reranker')
 
-    else:
-        edge_uuid_map = {}
-        search_result_uuids = []
+    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}
 
-        for result in search_results:
-            result_uuids = []
-            for edge in result:
-                result_uuids.append(edge.uuid)
-                edge_uuid_map[edge.uuid] = edge
+    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)
 
-            search_result_uuids.append(result_uuids)
+    reranked_nodes = [node_uuid_map[uuid] for uuid in reranked_uuids]
 
-        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
+    return reranked_nodes
 
-        reranked_uuids: list[str] = []
-        if config.reranker == Reranker.rrf:
-            reranked_uuids = rrf(search_result_uuids)
-        elif config.reranker == Reranker.node_distance:
-            if center_node_uuid is None:
-                logger.exception('No center node provided for Node Distance reranker')
-                raise Exception('No center node provided for Node Distance reranker')
-            reranked_uuids = await node_distance_reranker(
-                driver, search_result_uuids, center_node_uuid
-            )
-
-        reranked_edges = [edge_uuid_map[uuid] for uuid in reranked_uuids]
-        edges.extend(reranked_edges)
-
-    context = SearchResults(
-        episodes=episodes, nodes=nodes[: config.num_nodes], edges=edges[: config.num_edges]
-    )
 
-    end = time()
+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)
 
-    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
+    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 context
+    return reranked_communities

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/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,
+    )
+)