From c00def5b7122df2015ba0272f6f600ef770cd61d Mon Sep 17 00:00:00 2001
From: Zhichang Yu <yuzhichang@gmail.com>
Date: Fri, 14 Mar 2025 11:29:22 +0800
Subject: [PATCH] Fix 6030 (#6070)

### What problem does this PR solve?

Close #6030

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 graphrag/search.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/graphrag/search.py b/graphrag/search.py
index aa65fc6f5..e0df92235 100644
--- a/graphrag/search.py
+++ b/graphrag/search.py
@@ -19,6 +19,7 @@ from collections import defaultdict
 from copy import deepcopy
 import json_repair
 import pandas as pd
+import trio
 
 from api.utils import get_uuid
 from graphrag.query_analyze_prompt import PROMPTS
@@ -41,7 +42,7 @@ class KGSearch(Dealer):
         return response
 
     def query_rewrite(self, llm, question, idxnms, kb_ids):
-        ty2ents = get_entity_type2sampels(idxnms, kb_ids)
+        ty2ents = trio.run(lambda: get_entity_type2sampels(idxnms, kb_ids))
         hint_prompt = PROMPTS["minirag_query2kwd"].format(query=question,
                                                           TYPE_POOL=json.dumps(ty2ents, ensure_ascii=False, indent=2))
         result = self._chat(llm, hint_prompt, [{"role": "user", "content": "Output:"}], {"temperature": .5})