From f9de673d16bb0e1cc46733607993c41adb474de1 Mon Sep 17 00:00:00 2001
From: Alonso Guevara <alonsog@microsoft.com>
Date: Mon, 21 Oct 2024 15:34:33 -0600
Subject: [PATCH] Use asdict()

---
 graphrag/model/community_report.py            | 27 -------------------
 .../drift_search/drift_context.py             |  3 ++-
 2 files changed, 2 insertions(+), 28 deletions(-)

diff --git a/graphrag/model/community_report.py b/graphrag/model/community_report.py
index a16e448c..2666c0b5 100644
--- a/graphrag/model/community_report.py
+++ b/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,
-        }
diff --git a/graphrag/query/structured_search/drift_search/drift_context.py b/graphrag/query/structured_search/drift_search/drift_context.py
index 806cbd49..7911c54c 100644
--- a/graphrag/query/structured_search/drift_search/drift_context.py
+++ b/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."