diff --git a/lightrag/operate.py b/lightrag/operate.py index 0ebcda07..fbd311dd 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -1562,13 +1562,14 @@ async def _find_most_related_edges_from_entities( query_param: QueryParam, knowledge_graph_inst: BaseGraphStorage, ): - all_related_edges = await asyncio.gather( - *[knowledge_graph_inst.get_node_edges(dp["entity_name"]) for dp in node_datas] - ) + node_names = [dp["entity_name"] for dp in node_datas] + batch_edges_dict = await knowledge_graph_inst.get_nodes_edges_batch(node_names) + all_edges = [] seen = set() - for this_edges in all_related_edges: + for node_name in node_names: + this_edges = batch_edges_dict.get(node_name, []) for e in this_edges: sorted_edge = tuple(sorted(e)) if sorted_edge not in seen: