Refactor graphrag to remove redis lock (#5828)

### What problem does this PR solve?

Refactor graphrag to remove redis lock

### Type of change

- [x] Refactoring

graphrag/general/extractor.py

@@ -93,7 +93,7 @@ class Extractor:
         return dict(maybe_nodes), dict(maybe_edges)
 
     async def __call__(
-        self, chunks: list[tuple[str, str]],
+        self, doc_id: str, chunks: list[str],
             callback: Callable | None = None
     ):
 
@@ -101,9 +101,9 @@ class Extractor:
         start_ts = trio.current_time()
         out_results = []
         async with trio.open_nursery() as nursery:
-            for i, (cid, ck) in enumerate(chunks):
+            for i, ck in enumerate(chunks):
                 ck = truncate(ck, int(self._llm.max_length*0.8))
-                nursery.start_soon(lambda: self._process_single_content((cid, ck), i, len(chunks), out_results))
+                nursery.start_soon(lambda: self._process_single_content((doc_id, ck), i, len(chunks), out_results))
 
         maybe_nodes = defaultdict(list)
         maybe_edges = defaultdict(list)
@@ -241,10 +241,13 @@ class Extractor:
     ) -> str:
         summary_max_tokens = 512
         use_description = truncate(description, summary_max_tokens)
+        description_list=use_description.split(GRAPH_FIELD_SEP),
+        if len(description_list) <= 12:
+            return use_description
         prompt_template = SUMMARIZE_DESCRIPTIONS_PROMPT
         context_base = dict(
             entity_name=entity_or_relation_name,
-            description_list=use_description.split(GRAPH_FIELD_SEP),
+            description_list=description_list,
             language=self._language,
         )
         use_prompt = prompt_template.format(**context_base)

graphrag/general/index.py

@@ -15,196 +15,353 @@
 #
 import json
 import logging
-from functools import reduce, partial
+from functools import partial
 import networkx as nx
 import trio
 
 from api import settings
+from graphrag.light.graph_extractor import GraphExtractor as LightKGExt
+from graphrag.general.graph_extractor import GraphExtractor as GeneralKGExt
 from graphrag.general.community_reports_extractor import CommunityReportsExtractor
 from graphrag.entity_resolution import EntityResolution
 from graphrag.general.extractor import Extractor
-from graphrag.general.graph_extractor import DEFAULT_ENTITY_TYPES
-from graphrag.utils import graph_merge, set_entity, get_relation, set_relation, get_entity, get_graph, set_graph, \
-    chunk_id, update_nodes_pagerank_nhop_neighbour
+from graphrag.utils import (
+    graph_merge,
+    set_entity,
+    get_relation,
+    set_relation,
+    get_entity,
+    get_graph,
+    set_graph,
+    chunk_id,
+    update_nodes_pagerank_nhop_neighbour,
+    does_graph_contains,
+    get_graph_doc_ids,
+)
 from rag.nlp import rag_tokenizer, search
-from rag.utils.redis_conn import RedisDistributedLock
+from rag.utils.redis_conn import REDIS_CONN
 
 
-class Dealer:
-    def __init__(self,
-                 extractor: Extractor,
-                 tenant_id: str,
-                 kb_id: str,
-                 llm_bdl,
-                 chunks: list[tuple[str, str]],
-                 language,
-                 entity_types=DEFAULT_ENTITY_TYPES,
-                 embed_bdl=None,
-                 callback=None
-                 ):
-        self.tenant_id = tenant_id
-        self.kb_id = kb_id
-        self.chunks = chunks
-        self.llm_bdl = llm_bdl
-        self.embed_bdl = embed_bdl
-        self.ext = extractor(self.llm_bdl, language=language,
-                        entity_types=entity_types,
-                        get_entity=partial(get_entity, tenant_id, kb_id),
-                        set_entity=partial(set_entity, tenant_id, kb_id, self.embed_bdl),
-                        get_relation=partial(get_relation, tenant_id, kb_id),
-                        set_relation=partial(set_relation, tenant_id, kb_id, self.embed_bdl)
-                        )
-        self.graph = nx.Graph()
-        self.callback = callback
+def graphrag_task_set(tenant_id, kb_id, doc_id) -> bool:
+    key = f"graphrag:{tenant_id}:{kb_id}"
+    ok = REDIS_CONN.set(key, doc_id, exp=3600 * 24)
+    if not ok:
+        raise Exception(f"Faild to set the {key} to {doc_id}")
 
-    async def __call__(self):
-        docids = list(set([docid for docid, _ in self.chunks]))
-        ents, rels = await self.ext(self.chunks, self.callback)
-        for en in ents:
-            self.graph.add_node(en["entity_name"], entity_type=en["entity_type"])#, description=en["description"])
 
-        for rel in rels:
-            self.graph.add_edge(
-                rel["src_id"],
-                rel["tgt_id"],
-                weight=rel["weight"],
-                #description=rel["description"]
+def graphrag_task_get(tenant_id, kb_id) -> str | None:
+    key = f"graphrag:{tenant_id}:{kb_id}"
+    doc_id = REDIS_CONN.get(key)
+    return doc_id
+
+
+async def run_graphrag(
+    row: dict,
+    language,
+    with_resolution: bool,
+    with_community: bool,
+    chat_model,
+    embedding_model,
+    callback,
+):
+    start = trio.current_time()
+    tenant_id, kb_id, doc_id = row["tenant_id"], str(row["kb_id"]), row["doc_id"]
+    chunks = []
+    for d in settings.retrievaler.chunk_list(
+        doc_id, tenant_id, [kb_id], fields=["content_with_weight", "doc_id"]
+    ):
+        chunks.append(d["content_with_weight"])
+
+    graph, doc_ids = await update_graph(
+        LightKGExt
+        if row["parser_config"]["graphrag"]["method"] != "general"
+        else GeneralKGExt,
+        tenant_id,
+        kb_id,
+        doc_id,
+        chunks,
+        language,
+        row["parser_config"]["graphrag"]["entity_types"],
+        chat_model,
+        embedding_model,
+        callback,
+    )
+    if not graph:
+        return
+    if with_resolution or with_community:
+        graphrag_task_set(tenant_id, kb_id, doc_id)
+    if with_resolution:
+        await resolve_entities(
+            graph,
+            doc_ids,
+            tenant_id,
+            kb_id,
+            doc_id,
+            chat_model,
+            embedding_model,
+            callback,
+        )
+    if with_community:
+        await extract_community(
+            graph,
+            doc_ids,
+            tenant_id,
+            kb_id,
+            doc_id,
+            chat_model,
+            embedding_model,
+            callback,
+        )
+    now = trio.current_time()
+    callback(msg=f"GraphRAG for doc {doc_id} done in {now - start:.2f} seconds.")
+    return
+
+
+async def update_graph(
+    extractor: Extractor,
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+    chunks: list[str],
+    language,
+    entity_types,
+    llm_bdl,
+    embed_bdl,
+    callback,
+):
+    contains = await does_graph_contains(tenant_id, kb_id, doc_id)
+    if contains:
+        callback(msg=f"Graph already contains {doc_id}, cancel myself")
+        return None, None
+    start = trio.current_time()
+    ext = extractor(
+        llm_bdl,
+        language=language,
+        entity_types=entity_types,
+        get_entity=partial(get_entity, tenant_id, kb_id),
+        set_entity=partial(set_entity, tenant_id, kb_id, embed_bdl),
+        get_relation=partial(get_relation, tenant_id, kb_id),
+        set_relation=partial(set_relation, tenant_id, kb_id, embed_bdl),
+    )
+    ents, rels = await ext(doc_id, chunks, callback)
+    subgraph = nx.Graph()
+    for en in ents:
+        subgraph.add_node(en["entity_name"], entity_type=en["entity_type"])
+
+    for rel in rels:
+        subgraph.add_edge(
+            rel["src_id"],
+            rel["tgt_id"],
+            weight=rel["weight"],
+            # description=rel["description"]
+        )
+    # TODO: infinity doesn't support array search
+    chunk = {
+        "content_with_weight": json.dumps(
+            nx.node_link_data(subgraph, edges="edges"), ensure_ascii=False, indent=2
+        ),
+        "knowledge_graph_kwd": "subgraph",
+        "kb_id": kb_id,
+        "source_id": [doc_id],
+        "available_int": 0,
+        "removed_kwd": "N",
+    }
+    cid = chunk_id(chunk)
+    await trio.to_thread.run_sync(
+        lambda: settings.docStoreConn.insert(
+            [{"id": cid, **chunk}], search.index_name(tenant_id), kb_id
+        )
+    )
+    now = trio.current_time()
+    callback(msg=f"generated subgraph for doc {doc_id} in {now - start:.2f} seconds.")
+    start = now
+
+    while True:
+        new_graph = subgraph
+        now_docids = set([doc_id])
+        old_graph, old_doc_ids = await get_graph(tenant_id, kb_id)
+        if old_graph is not None:
+            logging.info("Merge with an exiting graph...................")
+            new_graph = graph_merge(old_graph, subgraph)
+        await update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, new_graph, 2)
+        if old_doc_ids:
+            for old_doc_id in old_doc_ids:
+                now_docids.add(old_doc_id)
+        old_doc_ids2 = await get_graph_doc_ids(tenant_id, kb_id)
+        delta_doc_ids = set(old_doc_ids2) - set(old_doc_ids)
+        if delta_doc_ids:
+            callback(
+                msg="The global graph has changed during merging, try again"
             )
-
-        with RedisDistributedLock(self.kb_id, 60*60):
-            old_graph, old_doc_ids = get_graph(self.tenant_id, self.kb_id)
-            if old_graph is not None:
-                logging.info("Merge with an exiting graph...................")
-                self.graph = reduce(graph_merge, [old_graph, self.graph])
-            update_nodes_pagerank_nhop_neighbour(self.tenant_id, self.kb_id, self.graph, 2)
-            if old_doc_ids:
-                docids.extend(old_doc_ids)
-                docids = list(set(docids))
-            set_graph(self.tenant_id, self.kb_id, self.graph, docids)
+            await trio.sleep(1)
+            continue
+        break
+    await set_graph(tenant_id, kb_id, new_graph, list(now_docids))
+    now = trio.current_time()
+    callback(
+        msg=f"merging subgraph for doc {doc_id} into the global graph done in {now - start:.2f} seconds."
+    )
+    return new_graph, now_docids
 
 
-class WithResolution(Dealer):
-    def __init__(self,
-                 tenant_id: str,
-                 kb_id: str,
-                 llm_bdl,
-                 embed_bdl=None,
-                 callback=None
-                 ):
-        self.tenant_id = tenant_id
-        self.kb_id = kb_id
-        self.llm_bdl = llm_bdl
-        self.embed_bdl = embed_bdl
-        self.callback = callback
-    async def __call__(self):
-        with RedisDistributedLock(self.kb_id, 60*60):
-            self.graph, doc_ids = await trio.to_thread.run_sync(lambda: get_graph(self.tenant_id, self.kb_id))
-            if not self.graph:
-                logging.error(f"Faild to fetch the graph. tenant_id:{self.kb_id}, kb_id:{self.kb_id}")
-                if self.callback:
-                    self.callback(-1, msg="Faild to fetch the graph.")
-                return
+async def resolve_entities(
+    graph,
+    doc_ids,
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+    llm_bdl,
+    embed_bdl,
+    callback,
+):
+    working_doc_id = graphrag_task_get(tenant_id, kb_id)
+    if doc_id != working_doc_id:
+        callback(
+            msg=f"Another graphrag task of doc_id {working_doc_id} is working on this kb, cancel myself"
+        )
+        return
+    start = trio.current_time()
+    er = EntityResolution(
+        llm_bdl,
+        get_entity=partial(get_entity, tenant_id, kb_id),
+        set_entity=partial(set_entity, tenant_id, kb_id, embed_bdl),
+        get_relation=partial(get_relation, tenant_id, kb_id),
+        set_relation=partial(set_relation, tenant_id, kb_id, embed_bdl),
+    )
+    reso = await er(graph)
+    graph = reso.graph
+    callback(msg=f"Graph resolution removed {len(reso.removed_entities)} nodes.")
+    await update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, graph, 2)
+    callback(msg="Graph resolution updated pagerank.")
 
-            if self.callback:
-                self.callback(msg="Fetch the existing graph.")
-            er = EntityResolution(self.llm_bdl,
-                                  get_entity=partial(get_entity, self.tenant_id, self.kb_id),
-                                  set_entity=partial(set_entity, self.tenant_id, self.kb_id, self.embed_bdl),
-                                  get_relation=partial(get_relation, self.tenant_id, self.kb_id),
-                                  set_relation=partial(set_relation, self.tenant_id, self.kb_id, self.embed_bdl))
-            reso = await er(self.graph)
-            self.graph = reso.graph
-            logging.info("Graph resolution is done. Remove {} nodes.".format(len(reso.removed_entities)))
-            if self.callback:
-                self.callback(msg="Graph resolution is done. Remove {} nodes.".format(len(reso.removed_entities)))
-            await trio.to_thread.run_sync(lambda: update_nodes_pagerank_nhop_neighbour(self.tenant_id, self.kb_id, self.graph, 2))
-            await trio.to_thread.run_sync(lambda: set_graph(self.tenant_id, self.kb_id, self.graph, doc_ids))
+    working_doc_id = graphrag_task_get(tenant_id, kb_id)
+    if doc_id != working_doc_id:
+        callback(
+            msg=f"Another graphrag task of doc_id {working_doc_id} is working on this kb, cancel myself"
+        )
+        return
+    await set_graph(tenant_id, kb_id, graph, doc_ids)
 
-        await trio.to_thread.run_sync(lambda: settings.docStoreConn.delete({
-            "knowledge_graph_kwd": "relation",
-            "kb_id": self.kb_id,
-            "from_entity_kwd": reso.removed_entities
-        }, search.index_name(self.tenant_id), self.kb_id))
-        await trio.to_thread.run_sync(lambda: settings.docStoreConn.delete({
-            "knowledge_graph_kwd": "relation",
-            "kb_id": self.kb_id,
-            "to_entity_kwd": reso.removed_entities
-        }, search.index_name(self.tenant_id), self.kb_id))
-        await trio.to_thread.run_sync(lambda: settings.docStoreConn.delete({
-            "knowledge_graph_kwd": "entity",
-            "kb_id": self.kb_id,
-            "entity_kwd": reso.removed_entities
-        }, search.index_name(self.tenant_id), self.kb_id))
+    await trio.to_thread.run_sync(
+        lambda: settings.docStoreConn.delete(
+            {
+                "knowledge_graph_kwd": "relation",
+                "kb_id": kb_id,
+                "from_entity_kwd": reso.removed_entities,
+            },
+            search.index_name(tenant_id),
+            kb_id,
+        )
+    )
+    await trio.to_thread.run_sync(
+        lambda: settings.docStoreConn.delete(
+            {
+                "knowledge_graph_kwd": "relation",
+                "kb_id": kb_id,
+                "to_entity_kwd": reso.removed_entities,
+            },
+            search.index_name(tenant_id),
+            kb_id,
+        )
+    )
+    await trio.to_thread.run_sync(
+        lambda: settings.docStoreConn.delete(
+            {
+                "knowledge_graph_kwd": "entity",
+                "kb_id": kb_id,
+                "entity_kwd": reso.removed_entities,
+            },
+            search.index_name(tenant_id),
+            kb_id,
+        )
+    )
+    now = trio.current_time()
+    callback(msg=f"Graph resolution done in {now - start:.2f}s.")
 
 
-class WithCommunity(Dealer):
-    def __init__(self,
-                 tenant_id: str,
-                 kb_id: str,
-                 llm_bdl,
-                 embed_bdl=None,
-                 callback=None
-                 ):
+async def extract_community(
+    graph,
+    doc_ids,
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+    llm_bdl,
+    embed_bdl,
+    callback,
+):
+    working_doc_id = graphrag_task_get(tenant_id, kb_id)
+    if doc_id != working_doc_id:
+        callback(
+            msg=f"Another graphrag task of doc_id {working_doc_id} is working on this kb, cancel myself"
+        )
+        return
+    start = trio.current_time()
+    ext = CommunityReportsExtractor(
+        llm_bdl,
+        get_entity=partial(get_entity, tenant_id, kb_id),
+        set_entity=partial(set_entity, tenant_id, kb_id, embed_bdl),
+        get_relation=partial(get_relation, tenant_id, kb_id),
+        set_relation=partial(set_relation, tenant_id, kb_id, embed_bdl),
+    )
+    cr = await ext(graph, callback=callback)
+    community_structure = cr.structured_output
+    community_reports = cr.output
+    working_doc_id = graphrag_task_get(tenant_id, kb_id)
+    if doc_id != working_doc_id:
+        callback(
+            msg=f"Another graphrag task of doc_id {working_doc_id} is working on this kb, cancel myself"
+        )
+        return
+    await set_graph(tenant_id, kb_id, graph, doc_ids)
 
-        self.tenant_id = tenant_id
-        self.kb_id = kb_id
-        self.community_structure = None
-        self.community_reports = None
-        self.llm_bdl = llm_bdl
-        self.embed_bdl = embed_bdl
-        self.callback = callback
-    async def __call__(self):
-        with RedisDistributedLock(self.kb_id, 60*60):
-            self.graph, doc_ids = get_graph(self.tenant_id, self.kb_id)
-            if not self.graph:
-                logging.error(f"Faild to fetch the graph. tenant_id:{self.kb_id}, kb_id:{self.kb_id}")
-                if self.callback:
-                    self.callback(-1, msg="Faild to fetch the graph.")
-                return
-            if self.callback:
-                self.callback(msg="Fetch the existing graph.")
-
-            cr = CommunityReportsExtractor(self.llm_bdl,
-                                  get_entity=partial(get_entity, self.tenant_id, self.kb_id),
-                                  set_entity=partial(set_entity, self.tenant_id, self.kb_id, self.embed_bdl),
-                                  get_relation=partial(get_relation, self.tenant_id, self.kb_id),
-                                  set_relation=partial(set_relation, self.tenant_id, self.kb_id, self.embed_bdl))
-            cr = await cr(self.graph, callback=self.callback)
-            self.community_structure = cr.structured_output
-            self.community_reports = cr.output
-            await trio.to_thread.run_sync(lambda: set_graph(self.tenant_id, self.kb_id, self.graph, doc_ids))
-
-        if self.callback:
-            self.callback(msg="Graph community extraction is done. Indexing {} reports.".format(len(cr.structured_output)))
-
-        await trio.to_thread.run_sync(lambda: settings.docStoreConn.delete({
+    now = trio.current_time()
+    callback(
+        msg=f"Graph extracted {len(cr.structured_output)} communities in {now - start:.2f}s."
+    )
+    start = now
+    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,
+        )
+    )
+    for stru, rep in zip(community_structure, community_reports):
+        obj = {
+            "report": rep,
+            "evidences": "\n".join([f["explanation"] for f in stru["findings"]]),
+        }
+        chunk = {
+            "docnm_kwd": stru["title"],
+            "title_tks": rag_tokenizer.tokenize(stru["title"]),
+            "content_with_weight": json.dumps(obj, ensure_ascii=False),
+            "content_ltks": rag_tokenizer.tokenize(
+                obj["report"] + " " + obj["evidences"]
+            ),
             "knowledge_graph_kwd": "community_report",
-            "kb_id": self.kb_id
-        }, search.index_name(self.tenant_id), self.kb_id))
-
-        for stru, rep in zip(self.community_structure, self.community_reports):
-            obj = {
-                "report": rep,
-                "evidences": "\n".join([f["explanation"] for f in stru["findings"]])
-            }
-            chunk = {
-                "docnm_kwd": stru["title"],
-                "title_tks": rag_tokenizer.tokenize(stru["title"]),
-                "content_with_weight": json.dumps(obj, ensure_ascii=False),
-                "content_ltks": rag_tokenizer.tokenize(obj["report"] +" "+ obj["evidences"]),
-                "knowledge_graph_kwd": "community_report",
-                "weight_flt": stru["weight"],
-                "entities_kwd": stru["entities"],
-                "important_kwd": stru["entities"],
-                "kb_id": self.kb_id,
-                "source_id": doc_ids,
-                "available_int": 0
-            }
-            chunk["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(chunk["content_ltks"])
-            #try:
-            #    ebd, _ = self.embed_bdl.encode([", ".join(community["entities"])])
-            #    chunk["q_%d_vec" % len(ebd[0])] = ebd[0]
-            #except Exception as e:
-            #    logging.exception(f"Fail to embed entity relation: {e}")
-            await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(self.tenant_id)))
+            "weight_flt": stru["weight"],
+            "entities_kwd": stru["entities"],
+            "important_kwd": stru["entities"],
+            "kb_id": kb_id,
+            "source_id": doc_ids,
+            "available_int": 0,
+        }
+        chunk["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(
+            chunk["content_ltks"]
+        )
+        # try:
+        #    ebd, _ = embed_bdl.encode([", ".join(community["entities"])])
+        #    chunk["q_%d_vec" % len(ebd[0])] = ebd[0]
+        # except Exception as e:
+        #    logging.exception(f"Fail to embed entity relation: {e}")
+        await trio.to_thread.run_sync(
+            lambda: settings.docStoreConn.insert(
+                [{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id)
+            )
+        )
 
+    now = trio.current_time()
+    callback(
+        msg=f"Graph indexed {len(cr.structured_output)} communities in {now - start:.2f}s."
+    )
+    return community_structure, community_reports

graphrag/general/smoke.py

@@ -16,7 +16,7 @@
 
 import argparse
 import json
-
+import logging
 import networkx as nx
 import trio
 
@@ -26,42 +26,85 @@ from api.db.services.document_service import DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api.db.services.user_service import TenantService
-from graphrag.general.index import WithCommunity, Dealer, WithResolution
-from graphrag.light.graph_extractor import GraphExtractor
-from rag.utils.redis_conn import RedisDistributedLock
+from graphrag.general.graph_extractor import GraphExtractor
+from graphrag.general.index import update_graph, with_resolution, with_community
 
 settings.init_settings()
 
-if __name__ == "__main__":
+
+def callback(prog=None, msg="Processing..."):
+    logging.info(msg)
+
+
+async def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('-t', '--tenant_id', default=False, help="Tenant ID", action='store', required=True)
-    parser.add_argument('-d', '--doc_id', default=False, help="Document ID", action='store', required=True)
+    parser.add_argument(
+        "-t",
+        "--tenant_id",
+        default=False,
+        help="Tenant ID",
+        action="store",
+        required=True,
+    )
+    parser.add_argument(
+        "-d",
+        "--doc_id",
+        default=False,
+        help="Document ID",
+        action="store",
+        required=True,
+    )
     args = parser.parse_args()
     e, doc = DocumentService.get_by_id(args.doc_id)
     if not e:
         raise LookupError("Document not found.")
     kb_id = doc.kb_id
 
-    chunks = [d["content_with_weight"] for d in
-              settings.retrievaler.chunk_list(args.doc_id, args.tenant_id, [kb_id], max_count=6,
-                                              fields=["content_with_weight"])]
-    chunks = [("x", c) for c in chunks]
-
-    RedisDistributedLock.clean_lock(kb_id)
+    chunks = [
+        d["content_with_weight"]
+        for d in settings.retrievaler.chunk_list(
+            args.doc_id,
+            args.tenant_id,
+            [kb_id],
+            max_count=6,
+            fields=["content_with_weight"],
+        )
+    ]
 
     _, tenant = TenantService.get_by_id(args.tenant_id)
     llm_bdl = LLMBundle(args.tenant_id, LLMType.CHAT, tenant.llm_id)
     _, kb = KnowledgebaseService.get_by_id(kb_id)
     embed_bdl = LLMBundle(args.tenant_id, LLMType.EMBEDDING, kb.embd_id)
 
-    dealer = Dealer(GraphExtractor, args.tenant_id, kb_id, llm_bdl, chunks, "English", embed_bdl=embed_bdl)
-    trio.run(dealer())
-    print(json.dumps(nx.node_link_data(dealer.graph), ensure_ascii=False, indent=2))
+    graph, doc_ids = await update_graph(
+        GraphExtractor,
+        args.tenant_id,
+        kb_id,
+        args.doc_id,
+        chunks,
+        "English",
+        llm_bdl,
+        embed_bdl,
+        callback,
+    )
+    print(json.dumps(nx.node_link_data(graph), ensure_ascii=False, indent=2))
 
-    dealer = WithResolution(args.tenant_id, kb_id, llm_bdl, embed_bdl)
-    trio.run(dealer())
-    dealer = WithCommunity(args.tenant_id, kb_id, llm_bdl, embed_bdl)
-    trio.run(dealer())
+    await with_resolution(
+        args.tenant_id, kb_id, args.doc_id, llm_bdl, embed_bdl, callback
+    )
+    community_structure, community_reports = await with_community(
+        args.tenant_id, kb_id, args.doc_id, llm_bdl, embed_bdl, callback
+    )
 
-    print("------------------ COMMUNITY REPORT ----------------------\n", dealer.community_reports)
-    print(json.dumps(dealer.community_structure, ensure_ascii=False, indent=2))
+    print(
+        "------------------ COMMUNITY STRUCTURE--------------------\n",
+        json.dumps(community_structure, ensure_ascii=False, indent=2),
+    )
+    print(
+        "------------------ COMMUNITY REPORTS----------------------\n",
+        community_reports,
+    )
+
+
+if __name__ == "__main__":
+    trio.run(main)

graphrag/light/smoke.py

@@ -18,22 +18,42 @@ import argparse
 import json
 from api import settings
 import networkx as nx
+import logging
+import trio
 
 from api.db import LLMType
 from api.db.services.document_service import DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api.db.services.user_service import TenantService
-from graphrag.general.index import Dealer
+from graphrag.general.index import update_graph
 from graphrag.light.graph_extractor import GraphExtractor
-from rag.utils.redis_conn import RedisDistributedLock
 
 settings.init_settings()
 
-if __name__ == "__main__":
+
+def callback(prog=None, msg="Processing..."):
+    logging.info(msg)
+
+
+async def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('-t', '--tenant_id', default=False, help="Tenant ID", action='store', required=True)
-    parser.add_argument('-d', '--doc_id', default=False, help="Document ID", action='store', required=True)
+    parser.add_argument(
+        "-t",
+        "--tenant_id",
+        default=False,
+        help="Tenant ID",
+        action="store",
+        required=True,
+    )
+    parser.add_argument(
+        "-d",
+        "--doc_id",
+        default=False,
+        help="Document ID",
+        action="store",
+        required=True,
+    )
     args = parser.parse_args()
 
     e, doc = DocumentService.get_by_id(args.doc_id)
@@ -41,18 +61,36 @@ if __name__ == "__main__":
         raise LookupError("Document not found.")
     kb_id = doc.kb_id
 
-    chunks = [d["content_with_weight"] for d in
-              settings.retrievaler.chunk_list(args.doc_id, args.tenant_id, [kb_id], max_count=6,
-                                              fields=["content_with_weight"])]
-    chunks = [("x", c) for c in chunks]
-
-    RedisDistributedLock.clean_lock(kb_id)
+    chunks = [
+        d["content_with_weight"]
+        for d in settings.retrievaler.chunk_list(
+            args.doc_id,
+            args.tenant_id,
+            [kb_id],
+            max_count=6,
+            fields=["content_with_weight"],
+        )
+    ]
 
     _, tenant = TenantService.get_by_id(args.tenant_id)
     llm_bdl = LLMBundle(args.tenant_id, LLMType.CHAT, tenant.llm_id)
     _, kb = KnowledgebaseService.get_by_id(kb_id)
     embed_bdl = LLMBundle(args.tenant_id, LLMType.EMBEDDING, kb.embd_id)
 
-    dealer = Dealer(GraphExtractor, args.tenant_id, kb_id, llm_bdl, chunks, "English", embed_bdl=embed_bdl)
+    graph, doc_ids = await update_graph(
+        GraphExtractor,
+        args.tenant_id,
+        kb_id,
+        args.doc_id,
+        chunks,
+        "English",
+        llm_bdl,
+        embed_bdl,
+        callback,
+    )
 
-    print(json.dumps(nx.node_link_data(dealer.graph), ensure_ascii=False, indent=2))
+    print(json.dumps(nx.node_link_data(graph), ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    trio.run(main)

graphrag/utils.py

@@ -352,25 +352,57 @@ def set_relation(tenant_id, kb_id, embd_mdl, from_ent_name, to_ent_name, meta):
             chunk["q_%d_vec" % len(ebd)] = ebd
         settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id), kb_id)
 
+async def does_graph_contains(tenant_id, kb_id, doc_id):
+    # Get doc_ids of graph
+    fields = ["source_id"]
+    condition = {
+        "knowledge_graph_kwd": ["graph"],
+        "removed_kwd": "N",
+    }
+    res = await trio.to_thread.run_sync(lambda: settings.docStoreConn.search(fields, [], condition, [], OrderByExpr(), 0, 1, search.index_name(tenant_id), [kb_id]))
+    fields2 = settings.docStoreConn.getFields(res, fields)
+    graph_doc_ids = set()
+    for chunk_id in fields2.keys():
+        graph_doc_ids = set(fields2[chunk_id]["source_id"])
+    return doc_id in graph_doc_ids
 
-def get_graph(tenant_id, kb_id):
+async def get_graph_doc_ids(tenant_id, kb_id) -> list[str]:
+    conds = {
+        "fields": ["source_id"],
+        "removed_kwd": "N",
+        "size": 1,
+        "knowledge_graph_kwd": ["graph"]
+    }
+    res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search(conds, search.index_name(tenant_id), [kb_id]))
+    doc_ids = []
+    if res.total == 0:
+        return doc_ids
+    for id in res.ids:
+        doc_ids = res.field[id]["source_id"]
+    return doc_ids
+
+
+async def get_graph(tenant_id, kb_id):
     conds = {
         "fields": ["content_with_weight", "source_id"],
         "removed_kwd": "N",
         "size": 1,
         "knowledge_graph_kwd": ["graph"]
     }
-    res = settings.retrievaler.search(conds, search.index_name(tenant_id), [kb_id])
+    res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search(conds, search.index_name(tenant_id), [kb_id]))
+    if res.total == 0:
+        return None, []
     for id in res.ids:
         try:
             return json_graph.node_link_graph(json.loads(res.field[id]["content_with_weight"]), edges="edges"), \
                    res.field[id]["source_id"]
         except Exception:
             continue
-    return rebuild_graph(tenant_id, kb_id)
+    result = await rebuild_graph(tenant_id, kb_id)
+    return result
 
 
-def set_graph(tenant_id, kb_id, graph, docids):
+async def set_graph(tenant_id, kb_id, graph, docids):
     chunk = {
         "content_with_weight": json.dumps(nx.node_link_data(graph, edges="edges"), ensure_ascii=False,
                                           indent=2),
@@ -379,13 +411,13 @@ def set_graph(tenant_id, kb_id, graph, docids):
         "source_id": list(docids),
         "available_int": 0,
         "removed_kwd": "N"
-    }
-    res = settings.retrievaler.search({"knowledge_graph_kwd": "graph", "size": 1, "fields": []}, search.index_name(tenant_id), [kb_id])
+    }     
+    res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search({"knowledge_graph_kwd": "graph", "size": 1, "fields": []}, search.index_name(tenant_id), [kb_id]))
     if res.ids:
-        settings.docStoreConn.update({"knowledge_graph_kwd": "graph"}, chunk,
-                                     search.index_name(tenant_id), kb_id)
+        await trio.to_thread.run_sync(lambda: settings.docStoreConn.update({"knowledge_graph_kwd": "graph"}, chunk,
+                                     search.index_name(tenant_id), kb_id))
     else:
-        settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id), kb_id)
+        await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id), kb_id))
 
 
 def is_continuous_subsequence(subseq, seq):
@@ -430,7 +462,7 @@ def merge_tuples(list1, list2):
     return result
 
 
-def update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, graph, n_hop):
+async def update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, graph, n_hop):
     def n_neighbor(id):
         nonlocal graph, n_hop
         count = 0
@@ -460,10 +492,10 @@ def update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, graph, n_hop):
     for n, p in pr.items():
         graph.nodes[n]["pagerank"] = p
         try:
-            settings.docStoreConn.update({"entity_kwd": n, "kb_id": kb_id},
+            await trio.to_thread.run_sync(lambda: settings.docStoreConn.update({"entity_kwd": n, "kb_id": kb_id},
                                          {"rank_flt": p,
-                                          "n_hop_with_weight": json.dumps(n_neighbor(n), ensure_ascii=False)},
-                                         search.index_name(tenant_id), kb_id)
+                                          "n_hop_with_weight": json.dumps( (n), ensure_ascii=False)},
+                                         search.index_name(tenant_id), kb_id))
         except Exception as e:
             logging.exception(e)
 
@@ -480,21 +512,21 @@ def update_nodes_pagerank_nhop_neighbour(tenant_id, kb_id, graph, n_hop):
         "knowledge_graph_kwd": "ty2ents",
         "available_int": 0
     }
-    res = settings.retrievaler.search({"knowledge_graph_kwd": "ty2ents", "size": 1, "fields": []},
-                                      search.index_name(tenant_id), [kb_id])
+    res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search({"knowledge_graph_kwd": "ty2ents", "size": 1, "fields": []},
+                                      search.index_name(tenant_id), [kb_id]))
     if res.ids:
-        settings.docStoreConn.update({"knowledge_graph_kwd": "ty2ents"},
+        await trio.to_thread.run_sync(lambda: settings.docStoreConn.update({"knowledge_graph_kwd": "ty2ents"},
                                      chunk,
-                                     search.index_name(tenant_id), kb_id)
+                                     search.index_name(tenant_id), kb_id))
     else:
-        settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id), kb_id)
+        await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert([{"id": chunk_id(chunk), **chunk}], search.index_name(tenant_id), kb_id))
 
 
-def get_entity_type2sampels(idxnms, kb_ids: list):
-    es_res = settings.retrievaler.search({"knowledge_graph_kwd": "ty2ents", "kb_id": kb_ids,
+async def get_entity_type2sampels(idxnms, kb_ids: list):
+    es_res = await trio.to_thread.run_sync(lambda: settings.retrievaler.search({"knowledge_graph_kwd": "ty2ents", "kb_id": kb_ids,
                                        "size": 10000,
                                        "fields": ["content_with_weight"]},
-                                      idxnms, kb_ids)
+                                      idxnms, kb_ids))
 
     res = defaultdict(list)
     for id in es_res.ids:
@@ -522,18 +554,18 @@ def flat_uniq_list(arr, key):
     return list(set(res))
 
 
-def rebuild_graph(tenant_id, kb_id):
+async def rebuild_graph(tenant_id, kb_id):
     graph = nx.Graph()
     src_ids = []
     flds = ["entity_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", "weight_int", "knowledge_graph_kwd", "source_id"]
     bs = 256
     for i in range(0, 39*bs, bs):
-        es_res = settings.docStoreConn.search(flds, [],
+        es_res = await trio.to_thread.run_sync(lambda: settings.docStoreConn.search(flds, [],
                                  {"kb_id": kb_id, "knowledge_graph_kwd": ["entity", "relation"]},
                                  [],
                                  OrderByExpr(),
                                  i, bs, search.index_name(tenant_id), [kb_id]
-                                 )
+                                 ))
         tot = settings.docStoreConn.getTotal(es_res)
         if tot == 0:
             return None, None