Fix graph creation

.semversioner/next-release/patch-20250430001238698062.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Fix graph creation missing edge weights."
+}

graphrag/index/operations/finalize_entities.py

@@ -24,7 +24,7 @@ def finalize_entities(
     layout_enabled: bool = False,
 ) -> pd.DataFrame:
     """All the steps to transform final entities."""
-    graph = create_graph(relationships)
+    graph = create_graph(relationships, edge_attr=["weight"])
     graph_embeddings = None
     if embed_config is not None and embed_config.enabled:
         graph_embeddings = embed_graph(

graphrag/index/operations/finalize_relationships.py

@@ -19,7 +19,7 @@ def finalize_relationships(
     relationships: pd.DataFrame,
 ) -> pd.DataFrame:
     """All the steps to transform final relationships."""
-    graph = create_graph(relationships)
+    graph = create_graph(relationships, edge_attr=["weight"])
     degrees = compute_degree(graph)
 
     final_relationships = relationships.drop_duplicates(subset=["source", "target"])

graphrag/index/workflows/create_communities.py

@@ -52,7 +52,7 @@ def create_communities(
     seed: int | None = None,
 ) -> pd.DataFrame:
     """All the steps to transform final communities."""
-    graph = create_graph(relationships)
+    graph = create_graph(relationships, edge_attr=["weight"])
 
     clusters = cluster_graph(
         graph,

tests/verbs/test_finalize_graph.py

@@ -28,8 +28,8 @@ async def test_finalize_graph():
     nodes_actual = await load_table_from_storage("entities", context.storage)
     edges_actual = await load_table_from_storage("relationships", context.storage)
 
-    assert len(nodes_actual) == 251
-    assert len(edges_actual) == 372
+    assert len(nodes_actual) == 291
+    assert len(edges_actual) == 452
 
     # x and y will be zero with the default configuration, because we do not embed/umap
     assert nodes_actual["x"].sum() == 0
@@ -54,8 +54,8 @@ async def test_finalize_graph_umap():
     nodes_actual = await load_table_from_storage("entities", context.storage)
     edges_actual = await load_table_from_storage("relationships", context.storage)
 
-    assert len(nodes_actual) == 251
-    assert len(edges_actual) == 372
+    assert len(nodes_actual) == 291
+    assert len(edges_actual) == 452
 
     # x and y should have some value other than zero due to umap
     assert nodes_actual["x"].sum() != 0

tests/verbs/test_prune_graph.py

@@ -28,4 +28,4 @@ async def test_prune_graph():
 
     nodes_actual = await load_table_from_storage("entities", context.storage)
 
-    assert len(nodes_actual) == 20
+    assert len(nodes_actual) == 21