Refa:replace trio with asyncio (#11831)

### What problem does this PR solve?

change:
replace trio with asyncio

### Type of change
- [x] Refactoring

graphrag/general/community_reports_extractor.py

@@ -5,6 +5,7 @@ Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
 """
 
+import asyncio
 import logging
 import json
 import os
@@ -24,7 +25,6 @@ from graphrag.general.leiden import add_community_info2graph
 from rag.llm.chat_model import Base as CompletionLLM
 from graphrag.utils import perform_variable_replacements, dict_has_keys_with_types, chat_limiter
 from common.token_utils import num_tokens_from_string
-import trio
 
 
 @dataclass
@@ -101,14 +101,11 @@ class CommunityReportsExtractor(Extractor):
             text = perform_variable_replacements(self._extraction_prompt, variables=prompt_variables)
             async with chat_limiter:
                 try:
-                    with trio.move_on_after(180 if enable_timeout_assertion else 1000000000) as cancel_scope:
-                        if task_id and has_canceled(task_id):
-                            logging.info(f"Task {task_id} cancelled before LLM call.")
-                            raise TaskCanceledException(f"Task {task_id} was cancelled")
-                        response = await trio.to_thread.run_sync( self._chat, text, [{"role": "user", "content": "Output:"}], {}, task_id)
-                    if cancel_scope.cancelled_caught:
-                        logging.warning("extract_community_report._chat timeout, skipping...")
-                        return
+                    timeout = 180 if enable_timeout_assertion else 1000000000
+                    response = await asyncio.wait_for(asyncio.to_thread(self._chat,text,[{"role": "user", "content": "Output:"}],{},task_id),timeout=timeout)
+                except asyncio.TimeoutError:
+                    logging.warning("extract_community_report._chat timeout, skipping...")
+                    return
                 except Exception as e:
                     logging.error(f"extract_community_report._chat failed: {e}")
                     return
@@ -141,17 +138,25 @@ class CommunityReportsExtractor(Extractor):
             if callback:
                 callback(msg=f"Communities: {over}/{total}, used tokens: {token_count}")
 
-        st = trio.current_time()
-        async with trio.open_nursery() as nursery:
-            for level, comm in communities.items():
-                logging.info(f"Level {level}: Community: {len(comm.keys())}")
-                for community in comm.items():
-                    if task_id and has_canceled(task_id):
-                        logging.info(f"Task {task_id} cancelled before community processing.")
-                        raise TaskCanceledException(f"Task {task_id} was cancelled")
-                    nursery.start_soon(extract_community_report, community)
+        st = asyncio.get_running_loop().time()
+        tasks = []
+        for level, comm in communities.items():
+            logging.info(f"Level {level}: Community: {len(comm.keys())}")
+            for community in comm.items():
+                if task_id and has_canceled(task_id):
+                    logging.info(f"Task {task_id} cancelled before community processing.")
+                    raise TaskCanceledException(f"Task {task_id} was cancelled")
+                tasks.append(asyncio.create_task(extract_community_report(community)))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"Error in community processing: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
         if callback:
-            callback(msg=f"Community reports done in {trio.current_time() - st:.2f}s, used tokens: {token_count}")
+            callback(msg=f"Community reports done in {asyncio.get_running_loop().time() - st:.2f}s, used tokens: {token_count}")
 
         return CommunityReportsResult(
             structured_output=res_dict,

graphrag/general/extractor.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import asyncio
 import logging
 import os
 import re
@@ -21,7 +22,6 @@ from copy import deepcopy
 from typing import Callable
 
 import networkx as nx
-import trio
 
 from api.db.services.task_service import has_canceled
 from common.connection_utils import timeout
@@ -109,14 +109,14 @@ class Extractor:
 
     async def __call__(self, doc_id: str, chunks: list[str], callback: Callable | None = None, task_id: str = ""):
         self.callback = callback
-        start_ts = trio.current_time()
+        start_ts = asyncio.get_running_loop().time()
 
         async def extract_all(doc_id, chunks, max_concurrency=MAX_CONCURRENT_PROCESS_AND_EXTRACT_CHUNK, task_id=""):
             out_results = []
             error_count = 0
             max_errors = int(os.environ.get("GRAPHRAG_MAX_ERRORS", 3))
 
-            limiter = trio.Semaphore(max_concurrency)
+            limiter = asyncio.Semaphore(max_concurrency)
 
             async def worker(chunk_key_dp: tuple[str, str], idx: int, total: int, task_id=""):
                 nonlocal error_count
@@ -137,9 +137,19 @@ class Extractor:
                         if error_count > max_errors:
                             raise Exception(f"Maximum error count ({max_errors}) reached. Last errors: {str(e)}")
 
-            async with trio.open_nursery() as nursery:
-                for i, ck in enumerate(chunks):
-                    nursery.start_soon(worker, (doc_id, ck), i, len(chunks), task_id)
+            tasks = [
+                asyncio.create_task(worker((doc_id, ck), i, len(chunks), task_id))
+                for i, ck in enumerate(chunks)
+            ]
+
+            try:
+                await asyncio.gather(*tasks, return_exceptions=False)
+            except Exception as e:
+                logging.error(f"Error in worker: {str(e)}")
+                for t in tasks:
+                    t.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+                raise
 
             if error_count > 0:
                 warning_msg = f"Completed with {error_count} errors (out of {len(chunks)} chunks processed)"
@@ -166,7 +176,7 @@ class Extractor:
             for k, v in m_edges.items():
                 maybe_edges[tuple(sorted(k))].extend(v)
             sum_token_count += token_count
-        now = trio.current_time()
+        now = asyncio.get_running_loop().time()
         if self.callback:
             self.callback(msg=f"Entities and relationships extraction done, {len(maybe_nodes)} nodes, {len(maybe_edges)} edges, {sum_token_count} tokens, {now - start_ts:.2f}s.")
         start_ts = now
@@ -176,14 +186,23 @@ class Extractor:
         if task_id and has_canceled(task_id):
             raise TaskCanceledException(f"Task {task_id} was cancelled before nodes merging")
 
-        async with trio.open_nursery() as nursery:
-            for en_nm, ents in maybe_nodes.items():
-                nursery.start_soon(self._merge_nodes, en_nm, ents, all_entities_data, task_id)
+        tasks = [
+            asyncio.create_task(self._merge_nodes(en_nm, ents, all_entities_data, task_id))
+            for en_nm, ents in maybe_nodes.items()
+        ]
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"Error merging nodes: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
 
         if task_id and has_canceled(task_id):
             raise TaskCanceledException(f"Task {task_id} was cancelled after nodes merging")
 
-        now = trio.current_time()
+        now = asyncio.get_running_loop().time()
         if self.callback:
             self.callback(msg=f"Entities merging done, {now - start_ts:.2f}s.")
 
@@ -194,14 +213,26 @@ class Extractor:
         if task_id and has_canceled(task_id):
             raise TaskCanceledException(f"Task {task_id} was cancelled before relationships merging")
 
-        async with trio.open_nursery() as nursery:
-            for (src, tgt), rels in maybe_edges.items():
-                nursery.start_soon(self._merge_edges, src, tgt, rels, all_relationships_data, task_id)
+        tasks = []
+        for (src, tgt), rels in maybe_edges.items():
+            tasks.append(
+                asyncio.create_task(
+                    self._merge_edges(src, tgt, rels, all_relationships_data, task_id)
+                )
+            )
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"Error during relationships merging: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
 
         if task_id and has_canceled(task_id):
             raise TaskCanceledException(f"Task {task_id} was cancelled after relationships merging")
 
-        now = trio.current_time()
+        now = asyncio.get_running_loop().time()
         if self.callback:
             self.callback(msg=f"Relationships merging done, {now - start_ts:.2f}s.")
 
@@ -309,5 +340,5 @@ class Extractor:
             raise TaskCanceledException(f"Task {task_id} was cancelled during summary handling")
 
         async with chat_limiter:
-            summary = await trio.to_thread.run_sync(self._chat, "", [{"role": "user", "content": use_prompt}], {}, task_id)
+            summary = await asyncio.to_thread(self._chat, "", [{"role": "user", "content": use_prompt}], {}, task_id)
         return summary

graphrag/general/graph_extractor.py

@@ -5,11 +5,11 @@ Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
 """
 
+import asyncio
 import re
 from typing import Any
 from dataclasses import dataclass
 import tiktoken
-import trio
 
 from graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS
 from graphrag.general.graph_prompt import GRAPH_EXTRACTION_PROMPT, CONTINUE_PROMPT, LOOP_PROMPT
@@ -107,7 +107,7 @@ class GraphExtractor(Extractor):
         }
         hint_prompt = perform_variable_replacements(self._extraction_prompt, variables=variables)
         async with chat_limiter:
-            response = await trio.to_thread.run_sync(self._chat, hint_prompt, [{"role": "user", "content": "Output:"}], {}, task_id)
+            response = await asyncio.to_thread(self._chat,hint_prompt,[{"role": "user", "content": "Output:"}],{},task_id)
         token_count += num_tokens_from_string(hint_prompt + response)
 
         results = response or ""
@@ -117,7 +117,7 @@ class GraphExtractor(Extractor):
         for i in range(self._max_gleanings):
             history.append({"role": "user", "content": CONTINUE_PROMPT})
             async with chat_limiter:
-                response = await trio.to_thread.run_sync(lambda: self._chat("", history, {}))
+                response = await asyncio.to_thread(self._chat, "", history, {})
             token_count += num_tokens_from_string("\n".join([m["content"] for m in history]) + response)
             results += response or ""
 
@@ -127,7 +127,7 @@ class GraphExtractor(Extractor):
             history.append({"role": "assistant", "content": response})
             history.append({"role": "user", "content": LOOP_PROMPT})
             async with chat_limiter:
-                continuation = await trio.to_thread.run_sync(lambda: self._chat("", history))
+                continuation = await asyncio.to_thread(self._chat, "", history)
             token_count += num_tokens_from_string("\n".join([m["content"] for m in history]) + response)
             if continuation != "Y":
                 break

graphrag/general/index.py

@@ -13,12 +13,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import asyncio
 import json
 import logging
 import os
 
 import networkx as nx
-import trio
 
 from api.db.services.document_service import DocumentService
 from api.db.services.task_service import has_canceled
@@ -54,25 +54,35 @@ async def run_graphrag(
     callback,
 ):
     enable_timeout_assertion = os.environ.get("ENABLE_TIMEOUT_ASSERTION")
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     tenant_id, kb_id, doc_id = row["tenant_id"], str(row["kb_id"]), row["doc_id"]
     chunks = []
     for d in settings.retriever.chunk_list(doc_id, tenant_id, [kb_id], max_count=10000, fields=["content_with_weight", "doc_id"], sort_by_position=True):
         chunks.append(d["content_with_weight"])
 
-    with trio.fail_after(max(120, len(chunks) * 60 * 10) if enable_timeout_assertion else 10000000000):
-        subgraph = await generate_subgraph(
-            LightKGExt if "method" not in row["kb_parser_config"].get("graphrag", {}) or row["kb_parser_config"]["graphrag"]["method"] != "general" else GeneralKGExt,
-            tenant_id,
-            kb_id,
-            doc_id,
-            chunks,
-            language,
-            row["kb_parser_config"]["graphrag"].get("entity_types", []),
-            chat_model,
-            embedding_model,
-            callback,
+    timeout_sec = max(120, len(chunks) * 60 * 10) if enable_timeout_assertion else 10000000000
+
+    try:
+        subgraph = await asyncio.wait_for(
+            generate_subgraph(
+                LightKGExt if "method" not in row["kb_parser_config"].get("graphrag", {})
+                    or row["kb_parser_config"]["graphrag"]["method"] != "general"
+                else GeneralKGExt,
+                tenant_id,
+                kb_id,
+                doc_id,
+                chunks,
+                language,
+                row["kb_parser_config"]["graphrag"].get("entity_types", []),
+                chat_model,
+                embedding_model,
+                callback,
+            ),
+            timeout=timeout_sec,
         )
+    except asyncio.TimeoutError:
+        logging.error("generate_subgraph timeout")
+        raise
 
     if not subgraph:
         return
@@ -125,7 +135,7 @@ async def run_graphrag(
             )
     finally:
         graphrag_task_lock.release()
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"GraphRAG for doc {doc_id} done in {now - start:.2f} seconds.")
     return
 
@@ -145,7 +155,7 @@ async def run_graphrag_for_kb(
 ) -> dict:
     tenant_id, kb_id = row["tenant_id"], row["kb_id"]
     enable_timeout_assertion = os.environ.get("ENABLE_TIMEOUT_ASSERTION")
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     fields_for_chunks = ["content_with_weight", "doc_id"]
 
     if not doc_ids:
@@ -211,7 +221,7 @@ async def run_graphrag_for_kb(
         callback(msg=f"[GraphRAG] kb:{kb_id} has no available chunks in all documents, skip.")
         return {"ok_docs": [], "failed_docs": doc_ids, "total_docs": len(doc_ids), "total_chunks": 0, "seconds": 0.0}
 
-    semaphore = trio.Semaphore(max_parallel_docs)
+    semaphore = asyncio.Semaphore(max_parallel_docs)
 
     subgraphs: dict[str, object] = {}
     failed_docs: list[tuple[str, str]] = []  # (doc_id, error)
@@ -234,20 +244,28 @@ async def run_graphrag_for_kb(
             try:
                 msg = f"[GraphRAG] build_subgraph doc:{doc_id}"
                 callback(msg=f"{msg} start (chunks={len(chunks)}, timeout={deadline}s)")
-                with trio.fail_after(deadline):
-                    sg = await generate_subgraph(
-                        kg_extractor,
-                        tenant_id,
-                        kb_id,
-                        doc_id,
-                        chunks,
-                        language,
-                        kb_parser_config.get("graphrag", {}).get("entity_types", []),
-                        chat_model,
-                        embedding_model,
-                        callback,
-                        task_id=row["id"]
+
+                try:
+                    sg = await asyncio.wait_for(
+                        generate_subgraph(
+                            kg_extractor,
+                            tenant_id,
+                            kb_id,
+                            doc_id,
+                            chunks,
+                            language,
+                            kb_parser_config.get("graphrag", {}).get("entity_types", []),
+                            chat_model,
+                            embedding_model,
+                            callback,
+                            task_id=row["id"]
+                        ),
+                        timeout=deadline,
                     )
+                except asyncio.TimeoutError:
+                    failed_docs.append((doc_id, "timeout"))
+                    callback(msg=f"{msg} FAILED: timeout")
+                    return
                 if sg:
                     subgraphs[doc_id] = sg
                     callback(msg=f"{msg} done")
@@ -264,9 +282,15 @@ async def run_graphrag_for_kb(
         callback(msg=f"Task {row['id']} cancelled before processing documents.")
         raise TaskCanceledException(f"Task {row['id']} was cancelled")
 
-    async with trio.open_nursery() as nursery:
-        for doc_id in doc_ids:
-            nursery.start_soon(build_one, doc_id)
+    tasks = [asyncio.create_task(build_one(doc_id)) for doc_id in doc_ids]
+    try:
+        await asyncio.gather(*tasks, return_exceptions=False)
+    except Exception as e:
+        logging.error(f"Error in asyncio.gather: {e}")
+        for t in tasks:
+            t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise
 
     if has_canceled(row["id"]):
         callback(msg=f"Task {row['id']} cancelled after document processing.")
@@ -275,7 +299,7 @@ async def run_graphrag_for_kb(
     ok_docs = [d for d in doc_ids if d in subgraphs]
     if not ok_docs:
         callback(msg=f"[GraphRAG] kb:{kb_id} no subgraphs generated successfully, end.")
-        now = trio.current_time()
+        now = asyncio.get_running_loop().time()
         return {"ok_docs": [], "failed_docs": failed_docs, "total_docs": len(doc_ids), "total_chunks": total_chunks, "seconds": now - start}
 
     kb_lock = RedisDistributedLock(f"graphrag_task_{kb_id}", lock_value="batch_merge", timeout=1200)
@@ -313,7 +337,7 @@ async def run_graphrag_for_kb(
         kb_lock.release()
 
     if not with_resolution and not with_community:
-        now = trio.current_time()
+        now = asyncio.get_running_loop().time()
         callback(msg=f"[GraphRAG] KB merge done in {now - start:.2f}s. ok={len(ok_docs)} / total={len(doc_ids)}")
         return {"ok_docs": ok_docs, "failed_docs": failed_docs, "total_docs": len(doc_ids), "total_chunks": total_chunks, "seconds": now - start}
 
@@ -356,7 +380,7 @@ async def run_graphrag_for_kb(
     finally:
         kb_lock.release()
 
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"[GraphRAG] GraphRAG for KB {kb_id} done in {now - start:.2f} seconds. ok={len(ok_docs)} failed={len(failed_docs)} total_docs={len(doc_ids)} total_chunks={total_chunks}")
     return {
         "ok_docs": ok_docs,
@@ -388,7 +412,7 @@ async def generate_subgraph(
     if contains:
         callback(msg=f"Graph already contains {doc_id}")
         return None
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     ext = extractor(
         llm_bdl,
         language=language,
@@ -436,9 +460,9 @@ async def generate_subgraph(
         "removed_kwd": "N",
     }
     cid = chunk_id(chunk)
-    await trio.to_thread.run_sync(settings.docStoreConn.delete, {"knowledge_graph_kwd": "subgraph", "source_id": doc_id}, search.index_name(tenant_id), kb_id)
-    await trio.to_thread.run_sync(settings.docStoreConn.insert, [{"id": cid, **chunk}], search.index_name(tenant_id), kb_id)
-    now = trio.current_time()
+    await asyncio.to_thread(settings.docStoreConn.delete,{"knowledge_graph_kwd": "subgraph", "source_id": doc_id},search.index_name(tenant_id),kb_id,)
+    await asyncio.to_thread(settings.docStoreConn.insert,[{"id": cid, **chunk}],search.index_name(tenant_id),kb_id,)
+    now = asyncio.get_running_loop().time()
     callback(msg=f"generated subgraph for doc {doc_id} in {now - start:.2f} seconds.")
     return subgraph
 
@@ -452,7 +476,7 @@ async def merge_subgraph(
     embedding_model,
     callback,
 ):
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     change = GraphChange()
     old_graph = await get_graph(tenant_id, kb_id, subgraph.graph["source_id"])
     if old_graph is not None:
@@ -468,7 +492,7 @@ async def merge_subgraph(
         new_graph.nodes[node_name]["pagerank"] = pagerank
 
     await set_graph(tenant_id, kb_id, embedding_model, new_graph, change, callback)
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"merging subgraph for doc {doc_id} into the global graph done in {now - start:.2f} seconds.")
     return new_graph
 
@@ -490,7 +514,7 @@ async def resolve_entities(
         callback(msg=f"Task {task_id} cancelled during entity resolution.")
         raise TaskCanceledException(f"Task {task_id} was cancelled")
 
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     er = EntityResolution(
         llm_bdl,
     )
@@ -505,7 +529,7 @@ async def resolve_entities(
         raise TaskCanceledException(f"Task {task_id} was cancelled")
 
     await set_graph(tenant_id, kb_id, embed_bdl, graph, change, callback)
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"Graph resolution done in {now - start:.2f}s.")
 
 
@@ -524,7 +548,7 @@ async def extract_community(
         callback(msg=f"Task {task_id} cancelled before community extraction.")
         raise TaskCanceledException(f"Task {task_id} was cancelled")
 
-    start = trio.current_time()
+    start = asyncio.get_running_loop().time()
     ext = CommunityReportsExtractor(
         llm_bdl,
     )
@@ -538,7 +562,7 @@ async def extract_community(
     community_reports = cr.output
     doc_ids = graph.graph["source_id"]
 
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"Graph extracted {len(cr.structured_output)} communities in {now - start:.2f}s.")
     start = now
     if task_id and has_canceled(task_id):
@@ -568,16 +592,10 @@ async def extract_community(
         chunk["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(chunk["content_ltks"])
         chunks.append(chunk)
 
-    await trio.to_thread.run_sync(
-        lambda: settings.docStoreConn.delete(
-            {"knowledge_graph_kwd": "community_report", "kb_id": kb_id},
-            search.index_name(tenant_id),
-            kb_id,
-        )
-    )
+    await asyncio.to_thread(settings.docStoreConn.delete,{"knowledge_graph_kwd": "community_report", "kb_id": kb_id},search.index_name(tenant_id),kb_id,)
     es_bulk_size = 4
     for b in range(0, len(chunks), es_bulk_size):
-        doc_store_result = await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert(chunks[b : b + es_bulk_size], search.index_name(tenant_id), kb_id))
+        doc_store_result = await asyncio.to_thread(settings.docStoreConn.insert,chunks[b : b + es_bulk_size],search.index_name(tenant_id),kb_id,)
         if doc_store_result:
             error_message = f"Insert chunk error: {doc_store_result}, please check log file and Elasticsearch/Infinity status!"
             raise Exception(error_message)
@@ -586,6 +604,6 @@ async def extract_community(
         callback(msg=f"Task {task_id} cancelled after community indexing.")
         raise TaskCanceledException(f"Task {task_id} was cancelled")
 
-    now = trio.current_time()
+    now = asyncio.get_running_loop().time()
     callback(msg=f"Graph indexed {len(cr.structured_output)} communities in {now - start:.2f}s.")
     return community_structure, community_reports

graphrag/general/mind_map_extractor.py

@@ -14,12 +14,12 @@
 #  limitations under the License.
 #
 
+import asyncio
 import logging
 import collections
 import re
 from typing import Any
 from dataclasses import dataclass
-import trio
 
 from graphrag.general.extractor import Extractor
 from graphrag.general.mind_map_prompt import MIND_MAP_EXTRACTION_PROMPT
@@ -89,17 +89,30 @@ class MindMapExtractor(Extractor):
         token_count = max(self._llm.max_length * 0.8, self._llm.max_length - 512)
         texts = []
         cnt = 0
-        async with trio.open_nursery() as nursery:
-            for i in range(len(sections)):
-                section_cnt = num_tokens_from_string(sections[i])
-                if cnt + section_cnt >= token_count and texts:
-                    nursery.start_soon(self._process_document, "".join(texts), prompt_variables, res)
-                    texts = []
-                    cnt = 0
-                texts.append(sections[i])
-                cnt += section_cnt
-            if texts:
-                nursery.start_soon(self._process_document, "".join(texts), prompt_variables, res)
+        tasks = []
+        for i in range(len(sections)):
+            section_cnt = num_tokens_from_string(sections[i])
+            if cnt + section_cnt >= token_count and texts:
+                tasks.append(asyncio.create_task(
+                    self._process_document("".join(texts), prompt_variables, res)
+                ))
+                texts = []
+                cnt = 0
+
+            texts.append(sections[i])
+            cnt += section_cnt
+        if texts:
+            tasks.append(asyncio.create_task(
+                self._process_document("".join(texts), prompt_variables, res)
+            ))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"Error processing document: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
         if not res:
             return MindMapResult(output={"id": "root", "children": []})
         merge_json = reduce(self._merge, res)
@@ -172,7 +185,7 @@ class MindMapExtractor(Extractor):
         }
         text = perform_variable_replacements(self._mind_map_prompt, variables=variables)
         async with chat_limiter:
-            response = await trio.to_thread.run_sync(lambda: self._chat(text, [{"role": "user", "content": "Output:"}], {}))
+            response = await asyncio.to_thread(self._chat,text,[{"role": "user", "content": "Output:"}],{})
         response = re.sub(r"```[^\n]*", "", response)
         logging.debug(response)
         logging.debug(self._todict(markdown_to_json.dictify(response)))

graphrag/general/smoke.py

@@ -15,10 +15,10 @@
 #
 
 import argparse
+import asyncio
 import json
 import logging
 import networkx as nx
-import trio
 
 from common.constants import LLMType
 from api.db.services.document_service import DocumentService
@@ -107,4 +107,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    trio.run(main)
+    asyncio.run(main)