fix: respect search_methods configuration in edge_search and node_search The search functions were ignoring the configured search_methods and always executing all three methods (BM25, cosine similarity, and BFS). This was causing unnecessary performance overhead - running 3x the required computations when users only wanted specific search methods. Now the functions build search tasks dynamically based on what's actually configured in search_methods, only executing the requested search algorithms. This should significantly improve performance when users need only specific search methods. Would appreciate your review on this approach\! Fixes #788
This commit is contained in:
parent
708d30cd6a
commit
506bc4fc98
1 changed files with 63 additions and 41 deletions
|
|
@ -177,31 +177,42 @@ async def edge_search(
|
|||
) -> tuple[list[EntityEdge], list[float]]:
|
||||
if config is None:
|
||||
return [], []
|
||||
search_results: list[list[EntityEdge]] = list(
|
||||
await semaphore_gather(
|
||||
*[
|
||||
edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
|
||||
edge_similarity_search(
|
||||
driver,
|
||||
query_vector,
|
||||
None,
|
||||
None,
|
||||
search_filter,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
config.sim_min_score,
|
||||
),
|
||||
edge_bfs_search(
|
||||
driver,
|
||||
bfs_origin_node_uuids,
|
||||
config.bfs_max_depth,
|
||||
search_filter,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
),
|
||||
]
|
||||
|
||||
# Build search tasks based on configured search methods
|
||||
search_tasks = []
|
||||
if EdgeSearchMethod.bm25 in config.search_methods:
|
||||
search_tasks.append(
|
||||
edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
|
||||
)
|
||||
)
|
||||
if EdgeSearchMethod.cosine_similarity in config.search_methods:
|
||||
search_tasks.append(
|
||||
edge_similarity_search(
|
||||
driver,
|
||||
query_vector,
|
||||
None,
|
||||
None,
|
||||
search_filter,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
config.sim_min_score,
|
||||
)
|
||||
)
|
||||
if EdgeSearchMethod.bfs in config.search_methods:
|
||||
search_tasks.append(
|
||||
edge_bfs_search(
|
||||
driver,
|
||||
bfs_origin_node_uuids,
|
||||
config.bfs_max_depth,
|
||||
search_filter,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
)
|
||||
)
|
||||
|
||||
# Execute only the configured search methods
|
||||
search_results: list[list[EntityEdge]] = []
|
||||
if search_tasks:
|
||||
search_results = list(await semaphore_gather(*search_tasks))
|
||||
|
||||
if EdgeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
|
||||
source_node_uuids = [edge.source_node_uuid for result in search_results for edge in result]
|
||||
|
|
@ -289,24 +300,35 @@ async def node_search(
|
|||
) -> tuple[list[EntityNode], list[float]]:
|
||||
if config is None:
|
||||
return [], []
|
||||
search_results: list[list[EntityNode]] = list(
|
||||
await semaphore_gather(
|
||||
*[
|
||||
node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
|
||||
node_similarity_search(
|
||||
driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
|
||||
),
|
||||
node_bfs_search(
|
||||
driver,
|
||||
bfs_origin_node_uuids,
|
||||
search_filter,
|
||||
config.bfs_max_depth,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
),
|
||||
]
|
||||
|
||||
# Build search tasks based on configured search methods
|
||||
search_tasks = []
|
||||
if NodeSearchMethod.bm25 in config.search_methods:
|
||||
search_tasks.append(
|
||||
node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
|
||||
)
|
||||
)
|
||||
if NodeSearchMethod.cosine_similarity in config.search_methods:
|
||||
search_tasks.append(
|
||||
node_similarity_search(
|
||||
driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
|
||||
)
|
||||
)
|
||||
if NodeSearchMethod.bfs in config.search_methods:
|
||||
search_tasks.append(
|
||||
node_bfs_search(
|
||||
driver,
|
||||
bfs_origin_node_uuids,
|
||||
search_filter,
|
||||
config.bfs_max_depth,
|
||||
group_ids,
|
||||
2 * limit,
|
||||
)
|
||||
)
|
||||
|
||||
# Execute only the configured search methods
|
||||
search_results: list[list[EntityNode]] = []
|
||||
if search_tasks:
|
||||
search_results = list(await semaphore_gather(*search_tasks))
|
||||
|
||||
if NodeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
|
||||
origin_node_uuids = [node.uuid for result in search_results for node in result]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue