Merge branch 'main' into fix__entity_name_key

microsoft · Nov 26, 2024 · 3e432f1 · 3e432f1
2 parents c8569e1 + ae796b9
commit 3e432f1
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 9 deletions.
diff --git a/.semversioner/next-release/patch-20241126205530514149.json b/.semversioner/next-release/patch-20241126205530514149.json
@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Fix Global Search with dynamic Community selection bug"
+}
diff --git a/graphrag/query/context_builder/dynamic_community_selection.py b/graphrag/query/context_builder/dynamic_community_selection.py
@@ -73,11 +73,12 @@ def __init__(
         }
         # mapping from level to communities
         self.levels: dict[str, list[str]] = {}
+
         for community in communities:
             if community.level not in self.levels:
                 self.levels[community.level] = []
-            if community.id in self.reports:
-                self.levels[community.level].append(community.id)
+            if community.short_id in self.reports:
+                self.levels[community.level].append(community.short_id)
 
         # start from root communities (level 0)
         self.starting_communities = self.levels["0"]
@@ -100,6 +101,7 @@ async def select(self, query: str) -> tuple[list[CommunityReport], dict[str, Any
             "output_tokens": 0,
         }
         relevant_communities = set()
+
         while queue:
             gather_results = await asyncio.gather(*[
                 rate_relevancy(

diff --git a/graphrag/query/factories.py b/graphrag/query/factories.py
@@ -3,8 +3,6 @@
 
 """Query Factory methods to support CLI."""
 
-from copy import deepcopy
-
 import tiktoken
 
 from graphrag.config.models.graph_rag_config import GraphRagConfig
@@ -105,12 +103,11 @@ def get_global_search_engine(
 
     dynamic_community_selection_kwargs = {}
     if dynamic_community_selection:
-        gs_config = config.global_search
-        _config = deepcopy(config)
-        _config.llm.model = _config.llm.deployment_name = gs_config.dynamic_search_llm
+        # TODO: Allow for another llm definition only for Global Search to leverage -mini models
+
         dynamic_community_selection_kwargs.update({
-            "llm": get_llm(_config),
-            "token_encoder": tiktoken.encoding_for_model(gs_config.dynamic_search_llm),
+            "llm": get_llm(config),
+            "token_encoder": tiktoken.encoding_for_model(config.llm.model),
             "keep_parent": gs_config.dynamic_search_keep_parent,
             "num_repeats": gs_config.dynamic_search_num_repeats,
             "use_summary": gs_config.dynamic_search_use_summary,