diff --git a/graphrag/entity_resolution.py b/graphrag/entity_resolution.py index a5be7226c..5d75027ea 100644 --- a/graphrag/entity_resolution.py +++ b/graphrag/entity_resolution.py @@ -111,7 +111,7 @@ class EntityResolution(Extractor): connect_graph.add_edges_from(resolution_result) async with trio.open_nursery() as nursery: for sub_connect_graph in nx.connected_components(connect_graph): - merging_nodes = list(sub_connect_graph.nodes()) + merging_nodes = list(sub_connect_graph) nursery.start_soon(lambda: self._merge_graph_nodes(graph, merging_nodes, change)) # Update pagerank diff --git a/graphrag/general/extractor.py b/graphrag/general/extractor.py index d4f9e9946..4c6f8e11f 100644 --- a/graphrag/general/extractor.py +++ b/graphrag/general/extractor.py @@ -201,7 +201,7 @@ class Extractor: # Merge two nodes, keep "entity_name", "entity_type", "page_rank" unchanged. node1_attrs = graph.nodes[node1] node0_attrs["description"] += f"{GRAPH_FIELD_SEP}{node1_attrs['description']}" - node0_attrs["source_id"] = sorted(set(node0_attrs["source_id"].extend(node1_attrs.get("source_id", [])))) + node0_attrs["source_id"] = sorted(set(node0_attrs["source_id"] + node1_attrs["source_id"])) for neighbor in graph.neighbors(node1): change.removed_edges.add(get_from_to(node1, neighbor)) if neighbor not in nodes_set: @@ -213,7 +213,7 @@ class Extractor: edge0_attrs["weight"] += edge1_attrs["weight"] edge0_attrs["description"] += f"{GRAPH_FIELD_SEP}{edge1_attrs['description']}" for attr in ["keywords", "source_id"]: - edge0_attrs[attr] = sorted(set(edge0_attrs[attr].extend(edge1_attrs.get(attr, [])))) + edge0_attrs[attr] = sorted(set(edge0_attrs[attr] + edge1_attrs[attr])) edge0_attrs["description"] = await self._handle_entity_relation_summary(f"({nodes[0]}, {neighbor})", edge0_attrs["description"]) graph.add_edge(nodes[0], neighbor, **edge0_attrs) else: