Fix basic search context and id mapping

Author: natoverse

graphrag/prompts/query/basic_search_system_prompt.py

@@ -27,7 +27,7 @@
 
 "Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Sources (2, 7, 64, 46, 34, +more)]. He is also CEO of company X [Data: Sources (1, 3)]"
 
-where 1, 2, 3, 7, 34, 46, and 64 represent the source id taken from the "source_id" column in the provided tables.
+where 1, 2, 3, 7, 34, 46, and 64 represent the source id taken from the "id" column in the provided tables.
 
 Do not include information where the supporting evidence for it is not provided.
 
@@ -60,7 +60,7 @@
 
 "Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Sources (2, 7, 64, 46, 34, +more)]. He is also CEO of company X [Data: Sources (1, 3)]"
 
-where 1, 2, 3, 7, 34, 46, and 64 represent the source id taken from the "source_id" column in the provided tables.
+where 1, 2, 3, 7, 34, 46, and 64 represent the source id taken from the "id" column in the provided tables.
 
 Do not include information where the supporting evidence for it is not provided.
 

graphrag/query/structured_search/basic_search/basic_context.py

@@ -38,7 +38,6 @@ def __init__(
         self.text_units = text_units
         self.text_unit_embeddings = text_unit_embeddings
         self.embedding_vectorstore_key = embedding_vectorstore_key
-        self.text_id_map = self._map_ids()
 
     def build_context(
         self,
@@ -48,7 +47,7 @@ def build_context(
         max_context_tokens: int = 12_000,
         context_name: str = "Sources",
         column_delimiter: str = "|",
-        text_id_col: str = "source_id",
+        text_id_col: str = "id",
         text_col: str = "text",
         **kwargs,
     ) -> ContextBuilderResult:
@@ -63,7 +62,7 @@ def build_context(
             text_unit_ids = {t.document.id for t in related_texts}
             text_units_filtered = []
             text_units_filtered = [
-                {text_id_col: t.id, text_col: t.text}
+                {text_id_col: t.short_id, text_col: t.text}
                 for t in self.text_units or []
                 if t.id in text_unit_ids
             ]
@@ -102,13 +101,5 @@ def build_context(
 
         return ContextBuilderResult(
             context_chunks=final_text,
-            context_records={context_name: final_text_df},
+            context_records={context_name.lower(): final_text_df},
         )
-
-    def _map_ids(self) -> dict[str, str]:
-        """Map id to short id in the text units."""
-        id_map = {}
-        text_units = self.text_units or []
-        for unit in text_units:
-            id_map[unit.id] = unit.short_id
-        return id_map