From 2bf7e7c018873d115cca2a67eef2eb719fcc6a31 Mon Sep 17 00:00:00 2001
From: Nathan Evans <github@talkswithnumbers.com>
Date: Thu, 18 Sep 2025 12:49:56 -0700
Subject: [PATCH] Fix multi-index search (#2063)

---
 .../patch-20250918192431890892.json           |  4 ++++
 graphrag/utils/api.py                         | 21 ++++++++++---------
 2 files changed, 15 insertions(+), 10 deletions(-)
 create mode 100644 .semversioner/next-release/patch-20250918192431890892.json

diff --git a/.semversioner/next-release/patch-20250918192431890892.json b/.semversioner/next-release/patch-20250918192431890892.json
new file mode 100644
index 00000000..88fde2b5
--- /dev/null
+++ b/.semversioner/next-release/patch-20250918192431890892.json
@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Fix multi-index search."
+}
diff --git a/graphrag/utils/api.py b/graphrag/utils/api.py
index 96ca15a6..09ac3bf2 100644
--- a/graphrag/utils/api.py
+++ b/graphrag/utils/api.py
@@ -168,57 +168,58 @@ def update_context_data(
     """
     updated_context_data = {}
     for key in context_data:
+        entries = context_data[key].to_dict(orient="records")
         updated_entry = []
         if key == "reports":
             updated_entry = [
                 dict(
-                    {k: entry[k] for k in entry},
+                    entry,
                     index_name=links["community_reports"][int(entry["id"])][
                         "index_name"
                     ],
                     index_id=links["community_reports"][int(entry["id"])]["id"],
                 )
-                for entry in context_data[key]
+                for entry in entries
             ]
         if key == "entities":
             updated_entry = [
                 dict(
-                    {k: entry[k] for k in entry},
+                    entry,
                     entity=entry["entity"].split("-")[0],
                     index_name=links["entities"][int(entry["id"])]["index_name"],
                     index_id=links["entities"][int(entry["id"])]["id"],
                 )
-                for entry in context_data[key]
+                for entry in entries
             ]
         if key == "relationships":
             updated_entry = [
                 dict(
-                    {k: entry[k] for k in entry},
+                    entry,
                     source=entry["source"].split("-")[0],
                     target=entry["target"].split("-")[0],
                     index_name=links["relationships"][int(entry["id"])]["index_name"],
                     index_id=links["relationships"][int(entry["id"])]["id"],
                 )
-                for entry in context_data[key]
+                for entry in entries
             ]
         if key == "claims":
             updated_entry = [
                 dict(
-                    {k: entry[k] for k in entry},
+                    entry,
                     entity=entry["entity"].split("-")[0],
                     index_name=links["covariates"][int(entry["id"])]["index_name"],
                     index_id=links["covariates"][int(entry["id"])]["id"],
                 )
-                for entry in context_data[key]
+                for entry in entries
             ]
         if key == "sources":
             updated_entry = [
                 dict(
-                    {k: entry[k] for k in entry},
+                    entry,
                     index_name=links["text_units"][int(entry["id"])]["index_name"],
                     index_id=links["text_units"][int(entry["id"])]["id"],
                 )
-                for entry in context_data[key]
+                for entry in entries
             ]
         updated_context_data[key] = updated_entry
     return updated_context_data