Use asdict()

graphrag/model/community_report.py

@@ -62,30 +62,3 @@ class CommunityReport(Named):
             full_content_embedding=d.get(full_content_embedding_key),
             attributes=d.get(attributes_key),
         )
-
-    def to_dict(
-        self,
-        id_key: str = "id",
-        title_key: str = "title",
-        community_id_key: str = "community_id",
-        short_id_key: str = "short_id",
-        summary_key: str = "summary",
-        full_content_key: str = "full_content",
-        rank_key: str = "rank",
-        summary_embedding_key: str = "summary_embedding",
-        full_content_embedding_key: str = "full_content_embedding",
-        attributes_key: str = "attributes",
-    ) -> dict[str, Any]:
-        """Convert the community report to a dictionary."""
-        return {
-            id_key: self.id,
-            title_key: self.title,
-            community_id_key: self.community_id,
-            short_id_key: self.short_id,
-            summary_key: self.summary,
-            full_content_key: self.full_content,
-            rank_key: self.rank,
-            summary_embedding_key: self.summary_embedding,
-            full_content_embedding_key: self.full_content_embedding,
-            attributes_key: self.attributes,
-        }

graphrag/query/structured_search/drift_search/drift_context.py

@@ -4,6 +4,7 @@
 """DRIFT Context Builder implementation."""
 
 import logging
+from dataclasses import asdict
 from typing import Any
 
 import numpy as np
@@ -111,7 +112,7 @@ class DRIFTSearchContextBuilder(DRIFTContextBuilder):
         ------
         ValueError: If some reports are missing full content or full content embeddings.
         """
-        report_df = pd.DataFrame([report.to_dict() for report in reports])
+        report_df = pd.DataFrame([asdict(report) for report in reports])
         missing_content_error = "Some reports are missing full content."
         missing_embedding_error = "Some reports are missing full content embeddings."