From 2d01e5d7b7d2e91c0835f7d53f88597c97f7230b Mon Sep 17 00:00:00 2001 From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com> Date: Mon, 26 Aug 2024 18:34:57 -0400 Subject: [PATCH] Search node centering (#45) * add new search reranker and update search * node distance reranking * format * rebase * no need for enumerate * mypy typing * defaultdict update * rrf prelim ranking --- graphiti_core/graphiti.py | 26 ++++++++++--- graphiti_core/search/search.py | 55 ++++++++++++++++++++-------- graphiti_core/search/search_utils.py | 42 ++++++++++++++++++++- 3 files changed, 101 insertions(+), 22 deletions(-) diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py index ec038f1e..6ff5b52b 100644 --- a/graphiti_core/graphiti.py +++ b/graphiti_core/graphiti.py @@ -26,7 +26,7 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge from graphiti_core.llm_client import LLMClient, OpenAIClient from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode -from graphiti_core.search.search import SearchConfig, hybrid_search +from graphiti_core.search.search import Reranker, SearchConfig, SearchMethod, hybrid_search from graphiti_core.search.search_utils import ( get_relevant_edges, get_relevant_nodes, @@ -515,7 +515,7 @@ async def add_episode_bulk( except Exception as e: raise e - async def search(self, query: str, num_results=10): + async def search(self, query: str, center_node_uuid: str | None = None, num_results=10): """ Perform a hybrid search on the knowledge graph. @@ -526,6 +526,8 @@ async def search(self, query: str, num_results=10): ---------- query : str The search query string. + center_node_uuid: str, optional + Facts will be reranked based on proximity to this node num_results : int, optional The maximum number of results to return. Defaults to 10. @@ -543,7 +545,14 @@ async def search(self, query: str, num_results=10): The search is performed using the current date and time as the reference point for temporal relevance. """ - search_config = SearchConfig(num_episodes=0, num_results=num_results) + 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, + search_methods=[SearchMethod.bm25, SearchMethod.cosine_similarity], + reranker=reranker, + ) edges = ( await hybrid_search( self.driver, @@ -551,6 +560,7 @@ async def search(self, query: str, num_results=10): query, datetime.now(), search_config, + center_node_uuid, ) ).edges @@ -558,7 +568,13 @@ async def search(self, query: str, num_results=10): return facts - async def _search(self, query: str, timestamp: datetime, config: SearchConfig): + async def _search( + self, + query: str, + timestamp: datetime, + config: SearchConfig, + center_node_uuid: str | None = None, + ): return await hybrid_search( - self.driver, self.llm_client.get_embedder(), query, timestamp, config + self.driver, self.llm_client.get_embedder(), query, timestamp, config, center_node_uuid ) diff --git a/graphiti_core/search/search.py b/graphiti_core/search/search.py index 956ae65d..03111225 100644 --- a/graphiti_core/search/search.py +++ b/graphiti_core/search/search.py @@ -16,6 +16,7 @@ import logging from datetime import datetime +from enum import Enum from time import time from neo4j import AsyncDriver @@ -28,6 +29,7 @@ edge_fulltext_search, edge_similarity_search, get_mentioned_nodes, + node_distance_reranker, rrf, ) from graphiti_core.utils import retrieve_episodes @@ -36,12 +38,22 @@ logger = logging.getLogger(__name__) +class SearchMethod(Enum): + cosine_similarity = 'cosine_similarity' + bm25 = 'bm25' + + +class Reranker(Enum): + rrf = 'reciprocal_rank_fusion' + node_distance = 'node_distance' + + class SearchConfig(BaseModel): - num_results: int = 10 + num_edges: int = 10 + num_nodes: int = 10 num_episodes: int = EPISODE_WINDOW_LEN - similarity_search: str = 'cosine' - text_search: str = 'BM25' - reranker: str = 'rrf' + search_methods: list[SearchMethod] + reranker: Reranker | None class SearchResults(BaseModel): @@ -51,7 +63,12 @@ class SearchResults(BaseModel): async def hybrid_search( - driver: AsyncDriver, embedder, query: str, timestamp: datetime, config: SearchConfig + driver: AsyncDriver, + embedder, + query: str, + timestamp: datetime, + config: SearchConfig, + center_node_uuid: str | None = None, ) -> SearchResults: start = time() @@ -65,11 +82,11 @@ async def hybrid_search( episodes.extend(await retrieve_episodes(driver, timestamp)) nodes.extend(await get_mentioned_nodes(driver, episodes)) - if config.text_search == 'BM25': + if SearchMethod.bm25 in config.search_methods: text_search = await edge_fulltext_search(query, driver) search_results.append(text_search) - if config.similarity_search == 'cosine': + if SearchMethod.cosine_similarity in config.search_methods: query_text = query.replace('\n', ' ') search_vector = ( (await embedder.create(input=[query_text], model='text-embedding-3-small')) @@ -80,19 +97,14 @@ async def hybrid_search( similarity_search = await edge_similarity_search(search_vector, driver) search_results.append(similarity_search) - if len(search_results) == 1: - edges = search_results[0] - - elif len(search_results) > 1 and config.reranker != 'rrf': + 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') - elif config.reranker == 'rrf': + else: edge_uuid_map = {} search_result_uuids = [] - logger.info([[edge.fact for edge in result] for result in search_results]) - for result in search_results: result_uuids = [] for edge in result: @@ -103,12 +115,23 @@ async def hybrid_search( search_result_uuids = [[edge.uuid for edge in result] for result in search_results] - reranked_uuids = rrf(search_result_uuids) + 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, edges=edges) + context = SearchResults( + episodes=episodes, nodes=nodes[: config.num_nodes], edges=edges[: config.num_edges] + ) end = time() diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index d73ea5e6..e9d658e0 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -333,7 +333,7 @@ async def get_relevant_edges( # takes in a list of rankings of uuids def rrf(results: list[list[str]], rank_const=1) -> list[str]: - scores: dict[str, int] = defaultdict(int) + scores: dict[str, float] = defaultdict(float) for result in results: for i, uuid in enumerate(result): scores[uuid] += 1 / (i + rank_const) @@ -344,3 +344,43 @@ def rrf(results: list[list[str]], rank_const=1) -> list[str]: sorted_uuids = [term[0] for term in scored_uuids] return sorted_uuids + + +async def node_distance_reranker( + driver: AsyncDriver, results: list[list[str]], center_node_uuid: str +) -> list[str]: + # use rrf as a preliminary ranker + sorted_uuids = rrf(results) + scores: dict[str, float] = {} + + for uuid in sorted_uuids: + # Find shortest path to center node + records, _, _ = await driver.execute_query( + """ + MATCH (source:Entity)-[r:RELATES_TO {uuid: $edge_uuid}]->(target:Entity) + MATCH p = SHORTEST 1 (center:Entity)-[:RELATES_TO]-+(n:Entity) + WHERE center.uuid = $center_uuid AND n.uuid IN [source.uuid, target.uuid] + RETURN min(length(p)) AS score, source.uuid AS source_uuid, target.uuid AS target_uuid + """, + edge_uuid=uuid, + center_uuid=center_node_uuid, + ) + distance = 0.01 + + for record in records: + if ( + record['source_uuid'] == center_node_uuid + or record['target_uuid'] == center_node_uuid + ): + continue + distance = record['score'] + + if uuid in scores: + scores[uuid] = min(1 / distance, scores[uuid]) + else: + scores[uuid] = 1 / distance + + # rerank on shortest distance + sorted_uuids.sort(reverse=True, key=lambda cur_uuid: scores[cur_uuid]) + + return sorted_uuids