Skip to content

Commit ae796b9

Browse files
Fix dynamic community selection in global search (#1450)
* Fix dynamic community selection in global search * Format * Ruff fix
1 parent 6d21ef2 commit ae796b9

File tree

3 files changed

+12
-9
lines changed

3 files changed

+12
-9
lines changed
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "patch",
3+
"description": "Fix Global Search with dynamic Community selection bug"
4+
}

graphrag/query/context_builder/dynamic_community_selection.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,12 @@ def __init__(
7373
}
7474
# mapping from level to communities
7575
self.levels: dict[str, list[str]] = {}
76+
7677
for community in communities:
7778
if community.level not in self.levels:
7879
self.levels[community.level] = []
79-
if community.id in self.reports:
80-
self.levels[community.level].append(community.id)
80+
if community.short_id in self.reports:
81+
self.levels[community.level].append(community.short_id)
8182

8283
# start from root communities (level 0)
8384
self.starting_communities = self.levels["0"]
@@ -100,6 +101,7 @@ async def select(self, query: str) -> tuple[list[CommunityReport], dict[str, Any
100101
"output_tokens": 0,
101102
}
102103
relevant_communities = set()
104+
103105
while queue:
104106
gather_results = await asyncio.gather(*[
105107
rate_relevancy(

graphrag/query/factories.py

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@
33

44
"""Query Factory methods to support CLI."""
55

6-
from copy import deepcopy
7-
86
import tiktoken
97

108
from graphrag.config.models.graph_rag_config import GraphRagConfig
@@ -105,12 +103,11 @@ def get_global_search_engine(
105103

106104
dynamic_community_selection_kwargs = {}
107105
if dynamic_community_selection:
108-
gs_config = config.global_search
109-
_config = deepcopy(config)
110-
_config.llm.model = _config.llm.deployment_name = gs_config.dynamic_search_llm
106+
# TODO: Allow for another llm definition only for Global Search to leverage -mini models
107+
111108
dynamic_community_selection_kwargs.update({
112-
"llm": get_llm(_config),
113-
"token_encoder": tiktoken.encoding_for_model(gs_config.dynamic_search_llm),
109+
"llm": get_llm(config),
110+
"token_encoder": tiktoken.encoding_for_model(config.llm.model),
114111
"keep_parent": gs_config.dynamic_search_keep_parent,
115112
"num_repeats": gs_config.dynamic_search_num_repeats,
116113
"use_summary": gs_config.dynamic_search_use_summary,

0 commit comments

Comments
 (0)