fix(graphrag): align llm typing and mind map annotations

What problem does this solve?

Enabling runtime type checking with beartype against the current RAGFlow GraphRAG code surfaced several latent mismatches between type annotations and actual runtime behavior.

The failures fall into three groups:
- GraphRAG components annotate their LLM dependency as the concrete chat model Base, but callers may pass compatible wrappers such as LLMBundle.
- MindMapExtractor contains several narrow or incorrect annotations that do not match the actual structures returned by markdown_to_json.dictify(), especially list and string leaves.
- general/index.py contains signatures that are stricter than real KB-level call sites, where extractor is passed as a class and doc_id can be None.

Without these fixes, runtime type checking can cause upload_and_parse and related GraphRAG flows to fail before or during mind map / graph processing.

Type of change

- Bug fix
- Type annotation correction
- Test coverage improvement

Why is this change recommended?

These issues are latent in the current GraphRAG code and may remain unnoticed in default setups, but they represent real mismatches between declared contracts and actual runtime behavior.

This change:
- makes GraphRAG depend on the minimal async chat capability it actually uses;
- aligns function signatures with real call patterns;
- fixes MindMapExtractor annotations to match its real input and output shapes;
- preserves list-based mind map leaves correctly during post-processing.

How is it fixed?

1. Introduce rag.graphrag.llm_protocol.GraphRAGCompletionLLM as the minimal protocol GraphRAG actually needs:
- llm_name
- max_length
- async_chat(...)

2. Replace concrete Base annotations with this protocol across GraphRAG components:
- rag/graphrag/general/extractor.py
- rag/graphrag/general/mind_map_extractor.py
- rag/graphrag/general/community_reports_extractor.py
- rag/graphrag/general/graph_extractor.py
- rag/graphrag/light/graph_extractor.py
- rag/graphrag/entity_resolution.py
- rag/graphrag/search.py

3. Fix MindMapExtractor annotations and leaf handling:
- widen _todict() to support Mapping | list | str
- widen _be_children() to support Mapping | list[str] | str
- change _process_document() return type from str to None
- preserve plain string list leaves in _list_to_kv() instead of collapsing them into an empty dict

4. Normalize GraphRAG chat responses to support both return shapes currently present in the codebase.

RAGFlow commit 67937a6 fixed one concrete extraction failure by indexing response[0], which addresses the tuple-returning path. However, LLMBundle.async_chat in the current codebase returns plain text, so GraphRAG can also receive a string depending on the caller. This change makes GraphRAG explicitly handle both tuple[str, int] and str responses instead of depending on only one path.

5. Relax general/index.py signatures to match real usage:
- generate_subgraph(extractor: type[Extractor], ...)
- resolve_entities(doc_id: str | None, ...)
- extract_community(doc_id: str | None, ...)

Validation

- python3 -m py_compile rag/graphrag/general/index.py
- python3 -m compileall rag/graphrag
- PYTHONPATH=/Users/dxl/project/python/local/ragflow /Users/dxl/project/python/multirag/.venv/bin/python -m pytest test/unit_test/rag/graphrag/test_llm_protocol.py

The targeted regression tests cover:
- protocol-based LLM acceptance in MindMapExtractor
- tuple-based chat responses remaining supported
- _todict() preserving list leaves
- _be_children() accepting list leaves
- _process_document() returning None

Author: yuehong136

rag/graphrag/entity_resolution.py

@@ -24,10 +24,10 @@
 import networkx as nx
 
 from rag.graphrag.general.extractor import Extractor
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM
 from rag.nlp import is_english
 import editdistance
 from rag.graphrag.entity_resolution_prompt import ENTITY_RESOLUTION_PROMPT
-from rag.llm.chat_model import Base as CompletionLLM
 from rag.graphrag.utils import perform_variable_replacements, chat_limiter, GraphChange
 from api.db.services.task_service import has_canceled
 from common.exceptions import TaskCanceledException
@@ -57,7 +57,7 @@ class EntityResolution(Extractor):
 
     def __init__(
             self,
-            llm_invoker: CompletionLLM,
+            llm_invoker: GraphRAGCompletionLLM,
     ):
         super().__init__(llm_invoker)
         """Init method definition."""
@@ -294,4 +294,3 @@ def is_similarity(self, a, b):
             return len(a & b) > 1
 
         return len(a & b)*1./max_l >= 0.8
-

rag/graphrag/general/community_reports_extractor.py

@@ -22,10 +22,10 @@
 from common.exceptions import TaskCanceledException
 from common.connection_utils import timeout
 from rag.graphrag.general import leiden
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM
 from rag.graphrag.general.community_report_prompt import COMMUNITY_REPORT_PROMPT
 from rag.graphrag.general.extractor import Extractor
 from rag.graphrag.general.leiden import add_community_info2graph
-from rag.llm.chat_model import Base as CompletionLLM
 from rag.graphrag.utils import perform_variable_replacements, dict_has_keys_with_types, chat_limiter
 from common.token_utils import num_tokens_from_string
 
@@ -46,7 +46,7 @@ class CommunityReportsExtractor(Extractor):
 
     def __init__(
             self,
-            llm_invoker: CompletionLLM,
+            llm_invoker: GraphRAGCompletionLLM,
             max_report_length: int | None = None,
     ):
         super().__init__(llm_invoker)

rag/graphrag/general/extractor.py

@@ -26,6 +26,7 @@
 from api.db.services.task_service import has_canceled
 from common.connection_utils import timeout
 from common.token_utils import truncate
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM, unwrap_graphrag_chat_response
 from rag.graphrag.general.graph_prompt import SUMMARIZE_DESCRIPTIONS_PROMPT
 from rag.graphrag.utils import (
     GraphChange,
@@ -39,7 +40,6 @@
     split_string_by_multi_markers,
 )
 from common.misc_utils import thread_pool_exec
-from rag.llm.chat_model import Base as CompletionLLM
 from rag.prompts.generator import message_fit_in
 from common.exceptions import TaskCanceledException
 
@@ -50,11 +50,11 @@
 
 
 class Extractor:
-    _llm: CompletionLLM
+    _llm: GraphRAGCompletionLLM
 
     def __init__(
         self,
-        llm_invoker: CompletionLLM,
+        llm_invoker: GraphRAGCompletionLLM,
         language: str | None = "English",
         entity_types: list[str] | None = None,
     ):
@@ -78,7 +78,8 @@ def _chat(self, system, history, gen_conf={}, task_id=""):
                     raise TaskCanceledException(f"Task {task_id} was cancelled")
             try:
                 response = asyncio.run(self._llm.async_chat(system_msg[0]["content"], hist, conf))
-                response = re.sub(r"^.*</think>", "", response[0], flags=re.DOTALL)
+                response = unwrap_graphrag_chat_response(response)
+                response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
                 if response.find("**ERROR**") >= 0:
                     raise Exception(response)
                 set_llm_cache(self._llm.llm_name, system, response, history, gen_conf)

rag/graphrag/general/graph_extractor.py

@@ -14,9 +14,9 @@
 import tiktoken
 
 from rag.graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM
 from rag.graphrag.general.graph_prompt import GRAPH_EXTRACTION_PROMPT, CONTINUE_PROMPT, LOOP_PROMPT
 from rag.graphrag.utils import ErrorHandlerFn, perform_variable_replacements, chat_limiter, split_string_by_multi_markers
-from rag.llm.chat_model import Base as CompletionLLM
 import networkx as nx
 from common.token_utils import num_tokens_from_string
 
@@ -52,7 +52,7 @@ class GraphExtractor(Extractor):
 
     def __init__(
         self,
-        llm_invoker: CompletionLLM,
+        llm_invoker: GraphRAGCompletionLLM,
         language: str | None = "English",
         entity_types: list[str] | None = None,
         tuple_delimiter_key: str | None = None,

rag/graphrag/general/index.py

@@ -393,7 +393,7 @@ async def build_one(doc_id: str):
 
 
 async def generate_subgraph(
-    extractor: Extractor,
+    extractor: type[Extractor],
     tenant_id: str,
     kb_id: str,
     doc_id: str,
@@ -504,7 +504,7 @@ async def resolve_entities(
     subgraph_nodes: set[str],
     tenant_id: str,
     kb_id: str,
-    doc_id: str,
+    doc_id: str | None,
     llm_bdl,
     embed_bdl,
     callback,
@@ -539,7 +539,7 @@ async def extract_community(
     graph,
     tenant_id: str,
     kb_id: str,
-    doc_id: str,
+    doc_id: str | None,
     llm_bdl,
     embed_bdl,
     callback,

rag/graphrag/general/mind_map_extractor.py

@@ -18,13 +18,14 @@
 import logging
 import collections
 import re
+from collections.abc import Mapping
 from typing import Any
 from dataclasses import dataclass
 
 from rag.graphrag.general.extractor import Extractor
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM
 from rag.graphrag.general.mind_map_prompt import MIND_MAP_EXTRACTION_PROMPT
 from rag.graphrag.utils import ErrorHandlerFn, perform_variable_replacements, chat_limiter
-from rag.llm.chat_model import Base as CompletionLLM
 import markdown_to_json
 from functools import reduce
 from common.token_utils import num_tokens_from_string
@@ -44,7 +45,7 @@ class MindMapExtractor(Extractor):
 
     def __init__(
             self,
-            llm_invoker: CompletionLLM,
+            llm_invoker: GraphRAGCompletionLLM,
             prompt: str | None = None,
             input_text_key: str | None = None,
             on_error: ErrorHandlerFn | None = None,
@@ -59,7 +60,7 @@ def __init__(
     def _key(self, k):
         return re.sub(r"\*+", "", k)
 
-    def _be_children(self, obj: dict, keyset: set):
+    def _be_children(self, obj: Mapping[str, Any] | list[str] | str, keyset: set[str]) -> list[dict[str, Any]]:
         if isinstance(obj, str):
             obj = [obj]
         if isinstance(obj, list):
@@ -150,36 +151,35 @@ def _merge(self, d1, d2):
 
         return d2
 
-    def _list_to_kv(self, data):
+    def _list_to_kv(self, data: dict[str, Any]) -> dict[str, Any]:
         for key, value in data.items():
             if isinstance(value, dict):
                 self._list_to_kv(value)
             elif isinstance(value, list):
                 new_value = {}
+                has_nested_list = False
                 for i in range(len(value)):
                     if isinstance(value[i], list) and i > 0:
+                        has_nested_list = True
                         new_value[value[i - 1]] = value[i][0]
-                data[key] = new_value
+                data[key] = new_value if has_nested_list else value
             else:
                 continue
         return data
 
-    def _todict(self, layer: collections.OrderedDict):
-        to_ret = layer
-        if isinstance(layer, collections.OrderedDict):
+    def _todict(self, layer: Mapping[str, Any] | list[Any] | str) -> dict[str, Any] | list[Any] | str:
+        if isinstance(layer, collections.OrderedDict | dict):
             to_ret = dict(layer)
-
-        try:
             for key, value in to_ret.items():
                 to_ret[key] = self._todict(value)
-        except AttributeError:
-            pass
-
-        return self._list_to_kv(to_ret)
+            return self._list_to_kv(to_ret)
+        if isinstance(layer, list):
+            return [self._todict(value) for value in layer]
+        return layer
 
     async def _process_document(
-            self, text: str, prompt_variables: dict[str, str], out_res
-    ) -> str:
+            self, text: str, prompt_variables: dict[str, str], out_res: list[dict[str, Any] | list[Any] | str]
+    ) -> None:
         variables = {
             **prompt_variables,
             self._input_text_key: text,

rag/graphrag/light/graph_extractor.py

@@ -16,9 +16,9 @@
 import networkx as nx
 
 from rag.graphrag.general.extractor import ENTITY_EXTRACTION_MAX_GLEANINGS, Extractor
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM
 from rag.graphrag.light.graph_prompt import PROMPTS
 from rag.graphrag.utils import chat_limiter, pack_user_ass_to_openai_messages, split_string_by_multi_markers
-from rag.llm.chat_model import Base as CompletionLLM
 from common.token_utils import num_tokens_from_string
 
 @dataclass
@@ -34,7 +34,7 @@ class GraphExtractor(Extractor):
 
     def __init__(
         self,
-        llm_invoker: CompletionLLM,
+        llm_invoker: GraphRAGCompletionLLM,
         language: str | None = "English",
         entity_types: list[str] | None = None,
         example_number: int = 2,

rag/graphrag/llm_protocol.py

@@ -0,0 +1,25 @@
+from typing import Any, Protocol, TypeAlias, runtime_checkable
+
+GraphRAGChatResponse: TypeAlias = str | tuple[str, int]
+
+
+def unwrap_graphrag_chat_response(response: GraphRAGChatResponse) -> str:
+    if isinstance(response, tuple):
+        return response[0]
+    return response
+
+
+@runtime_checkable
+class GraphRAGCompletionLLM(Protocol):
+    """Minimal async chat contract used across GraphRAG components."""
+
+    llm_name: str
+    max_length: int
+
+    async def async_chat(
+        self,
+        system: str,
+        history: list[dict[str, Any]],
+        gen_conf: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> GraphRAGChatResponse: ...

rag/graphrag/search.py

@@ -22,6 +22,7 @@
 import pandas as pd
 
 from common.misc_utils import get_uuid
+from rag.graphrag.llm_protocol import GraphRAGCompletionLLM, unwrap_graphrag_chat_response
 from rag.graphrag.query_analyze_prompt import PROMPTS
 from rag.graphrag.utils import get_entity_type2samples, get_llm_cache, set_llm_cache, get_relation
 from common.token_utils import num_tokens_from_string
@@ -33,17 +34,18 @@
 
 
 class KGSearch(Dealer):
-    async def _chat(self, llm_bdl, system, history, gen_conf):
+    async def _chat(self, llm_bdl: GraphRAGCompletionLLM, system: str, history: list[dict[str, str]], gen_conf: dict):
         response = get_llm_cache(llm_bdl.llm_name, system, history, gen_conf)
         if response:
             return response
         response = await llm_bdl.async_chat(system, history, gen_conf)
+        response = unwrap_graphrag_chat_response(response)
         if response.find("**ERROR**") >= 0:
             raise Exception(response)
         set_llm_cache(llm_bdl.llm_name, system, response, history, gen_conf)
         return response
 
-    async def query_rewrite(self, llm, question, idxnms, kb_ids):
+    async def query_rewrite(self, llm: GraphRAGCompletionLLM, question, idxnms, kb_ids):
         ty2ents = await get_entity_type2samples(idxnms, kb_ids)
         hint_prompt = PROMPTS["minirag_query2kwd"].format(query=question,
                                                           TYPE_POOL=json.dumps(ty2ents, ensure_ascii=False, indent=2))

test/unit_test/rag/graphrag/test_llm_protocol.py

@@ -0,0 +1,99 @@
+import asyncio
+import collections
+import sys
+import types
+
+api_module = types.ModuleType("api")
+api_module.__path__ = []
+db_module = types.ModuleType("api.db")
+db_module.__path__ = []
+services_module = types.ModuleType("api.db.services")
+services_module.__path__ = []
+task_service_module = types.ModuleType("api.db.services.task_service")
+task_service_module.has_canceled = lambda *_args, **_kwargs: False
+
+api_module.db = db_module
+db_module.services = services_module
+services_module.task_service = task_service_module
+
+sys.modules.setdefault("api", api_module)
+sys.modules.setdefault("api.db", db_module)
+sys.modules.setdefault("api.db.services", services_module)
+sys.modules.setdefault("api.db.services.task_service", task_service_module)
+
+import rag.graphrag.general.extractor as extractor_module
+import rag.graphrag.general.mind_map_extractor as mind_map_extractor_module
+from rag.graphrag.general.mind_map_extractor import MindMapExtractor
+
+
+class FakeLLM:
+    llm_name = "fake-llm"
+    max_length = 4096
+
+    async def async_chat(self, system, history: list[dict[str, str]], gen_conf=None, **kwargs):
+        return "{}"
+
+
+class TupleLLM:
+    llm_name = "tuple-llm"
+    max_length = 4096
+
+    async def async_chat(self, system, history: list[dict[str, str]], gen_conf=None, **kwargs):
+        return "{}", 0
+
+
+def test_mind_map_extractor_accepts_protocol_based_llm():
+    extractor = MindMapExtractor(FakeLLM())
+
+    assert extractor._llm.llm_name == "fake-llm"
+    assert extractor._llm.max_length == 4096
+
+
+def test_mind_map_extractor_accepts_tuple_chat_response(monkeypatch):
+    extractor = MindMapExtractor(TupleLLM())
+    monkeypatch.setattr(extractor_module, "get_llm_cache", lambda *args, **kwargs: None)
+    monkeypatch.setattr(extractor_module, "set_llm_cache", lambda *args, **kwargs: None)
+
+    assert extractor._chat("system", [{"role": "user", "content": "Output:"}], {}) == "{}"
+
+
+def test_mind_map_extractor_todict_supports_list_leaves():
+    extractor = MindMapExtractor(FakeLLM())
+    layer = collections.OrderedDict(
+        {
+            "顶层": collections.OrderedDict(
+                {
+                    "部分A": [
+                        "点1",
+                        "点2",
+                    ]
+                }
+            )
+        }
+    )
+
+    assert extractor._todict(layer) == {"顶层": {"部分A": ["点1", "点2"]}}
+
+
+def test_mind_map_extractor_be_children_supports_list_leaves():
+    extractor = MindMapExtractor(FakeLLM())
+
+    assert extractor._be_children(["点1", "点2"], {"顶层"}) == [
+        {"id": "点1", "children": []},
+        {"id": "点2", "children": []},
+    ]
+
+
+def test_mind_map_extractor_process_document_returns_none(monkeypatch):
+    extractor = MindMapExtractor(FakeLLM())
+    out_res = []
+
+    async def fake_thread_pool_exec(*args, **kwargs):
+        return "# 顶层\n## 部分A\n- 点1\n- 点2"
+
+    monkeypatch.setattr(mind_map_extractor_module, "thread_pool_exec", fake_thread_pool_exec)
+
+    result = asyncio.run(extractor._process_document("课堂纪要", {}, out_res))
+
+    assert result is None
+    assert out_res == [{"顶层": {"部分A": ["点1", "点2"]}}]