Fix dynamic community selection in global search (#1450)

* Fix dynamic community selection in global search

* Format

* Ruff fix
This commit is contained in:
Alonso Guevara 2024-11-26 15:19:50 -06:00 committed by GitHub
parent 6d21ef2683
commit ae796b99cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 12 additions and 9 deletions

View File

@ -0,0 +1,4 @@
{
"type": "patch",
"description": "Fix Global Search with dynamic Community selection bug"
}

View File

@ -73,11 +73,12 @@ class DynamicCommunitySelection:
}
# 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 @@ class DynamicCommunitySelection:
"output_tokens": 0,
}
relevant_communities = set()
while queue:
gather_results = await asyncio.gather(*[
rate_relevancy(

View File

@ -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,