fix: fix operator partition_kg

Author: ChenZiHong-Gavin

graphgen/operators/partition/partition_kg.py

@@ -56,15 +56,7 @@ async def partition_kg(
     logger.info("Partitioned the graph into %d communities.", len(communities))
     batches = await partitioner.community2batch(communities, g=kg_instance)
 
-    for _, batch in enumerate(batches):
-        nodes, edges = batch
-        for node_id, node_data in nodes:
-            entity_type = node_data.get("entity_type")
-            if entity_type and "image" in entity_type.lower():
-                node_id = node_id.strip('"').lower()
-                image_data = await chunk_storage.get_by_id(node_id)
-                if image_data:
-                    node_data["images"] = image_data
+    batches = await attach_additional_data_to_node(batches, chunk_storage)
     return batches