feat: add bounded graph extraction concurrency

hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py

@@ -40,6 +40,7 @@ def extract_sync(req: GraphExtractRequest) -> GraphExtractResponse:
                 req.extract_type,
                 language=req.language,
                 split_type=req.split_type,
+                graph_extract_max_workers=req.graph_extract_max_workers,
                 client_config=req.client_config,
             )
             raw = json.loads(result_str)

hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py

@@ -21,6 +21,8 @@
 from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
+from hugegraph_llm.utils.graph_extract_config import MAX_GRAPH_EXTRACT_WORKERS
+
 
 class GraphExtractClientConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -44,6 +46,12 @@ class GraphExtractRequest(BaseModel):
     extract_type: Literal["property_graph"] = Query("property_graph", description="Extraction type.")
     language: Literal["zh", "en"] = Query("zh", description="Language for chunk splitting.")
     split_type: Literal["document", "paragraph", "sentence"] = Query("document", description="Chunk split granularity.")
+    graph_extract_max_workers: int = Query(
+        1,
+        ge=1,
+        le=MAX_GRAPH_EXTRACT_WORKERS,
+        description="Maximum parallel chunk extraction calls.",
+    )
     include_meta: bool = Query(False, description="Include vertex/edge/text counts in the response.")
     client_config: Optional[GraphExtractClientConfig] = Field(None, description="Request-scoped HugeGraph connection.")
 

hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py

@@ -79,6 +79,7 @@ class BasePromptConfig:
     gremlin_generate_prompt: str = ""
     doc_input_text: str = ""
     graph_extract_split_type: str = "document"
+    graph_extract_max_workers: int = 1
     schema_generator_query_examples: str = ""
     schema_generator_few_shot_examples: str = ""
     _language_generated: str = ""
@@ -140,6 +141,7 @@ def to_literal(val):
             "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt),
             "doc_input_text": to_literal(self.doc_input_text),
             "graph_extract_split_type": to_literal(self.graph_extract_split_type),
+            "graph_extract_max_workers": self.graph_extract_max_workers,
             "schema_generator_query_examples": to_literal(self.schema_generator_query_examples),
             "schema_generator_few_shot_examples": to_literal(self.schema_generator_few_shot_examples),
             "_language_generated": str(self.llm_settings.language).lower().strip(),

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py

@@ -153,17 +153,20 @@ def store_prompt(
     schema,
     example_prompt,
     graph_extract_split_type="document",
+    graph_extract_max_workers=1,
 ):
     if (
         prompt.doc_input_text != doc
         or prompt.graph_schema != schema
         or prompt.extract_graph_prompt != example_prompt
         or prompt.graph_extract_split_type != graph_extract_split_type
+        or prompt.graph_extract_max_workers != int(graph_extract_max_workers)
     ):
         prompt.doc_input_text = doc
         prompt.graph_schema = schema
         prompt.extract_graph_prompt = example_prompt
         prompt.graph_extract_split_type = graph_extract_split_type
+        prompt.graph_extract_max_workers = int(graph_extract_max_workers)
         prompt.update_yaml_file()
 
 
@@ -405,6 +408,14 @@ def create_vector_graph_block():
                 label="Graph Extraction Split Type",
                 info=("document keeps the current behavior; paragraph/sentence split long docs before extraction."),
             )
+            graph_extract_max_workers = gr.Slider(
+                minimum=1,
+                maximum=8,
+                value=prompt.graph_extract_max_workers,
+                step=1,
+                label="Graph Extraction Chunk Concurrency",
+                info="Maximum parallel chunk extraction calls. 1 keeps serial behavior.",
+            )
             graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary")
             graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True)
             graph_index_rebuild_bt = gr.Button("Update Vid Embedding")
@@ -440,6 +451,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         vector_index_btn1.click(clean_vector_index).then(
@@ -449,6 +461,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         vector_import_bt.click(build_vector_index, inputs=[input_file, input_text], outputs=out).then(
@@ -458,6 +471,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         graph_index_btn0.click(get_graph_index_info, outputs=out).then(
@@ -467,6 +481,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         graph_index_btn1.click(clean_all_graph_index).then(
@@ -476,6 +491,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         graph_data_btn0.click(clean_all_graph_data).then(
@@ -485,6 +501,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
         graph_index_rebuild_bt.click(update_vid_embedding, outputs=out).then(
@@ -494,6 +511,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
 
@@ -506,6 +524,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
             outputs=[out],
         ).then(
@@ -515,6 +534,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
 
@@ -527,6 +547,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],
         )
 
@@ -542,6 +563,7 @@ def create_vector_graph_block():
                 input_schema,
                 info_extract_template,
                 graph_split_type,
+                graph_extract_max_workers,
             ],  # Persist the updated schema-generator examples
         )
 

hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py

@@ -27,6 +27,7 @@
     VALID_SPLIT_TYPES,
 )
 from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState
+from hugegraph_llm.utils.graph_extract_config import validate_graph_extract_max_workers
 from hugegraph_llm.utils.log import log
 
 
@@ -44,6 +45,7 @@ def prepare(
         extract_type,
         split_type=SPLIT_TYPE_DOCUMENT,
         language="zh",
+        graph_extract_max_workers=1,
         **kwargs,
     ):
         # prepare input data
@@ -53,6 +55,7 @@ def prepare(
             raise ValueError("split_type must be document, paragraph, or sentence")
 
         prepared_input.split_type = split_type
+        prepared_input.graph_extract_max_workers = validate_graph_extract_max_workers(graph_extract_max_workers)
         prepared_input.example_prompt = example_prompt
         prepared_input.schema = schema
         prepared_input.extract_type = extract_type
@@ -76,6 +79,7 @@ def build_flow(
         extract_type,
         split_type=SPLIT_TYPE_DOCUMENT,
         language="zh",
+        graph_extract_max_workers=1,
         **kwargs,
     ):
         pipeline = GPipeline()
@@ -89,6 +93,7 @@ def build_flow(
             extract_type,
             split_type,
             language,
+            graph_extract_max_workers,
             **kwargs,
         )
 

hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py

@@ -41,7 +41,11 @@ def node_init(self):
         if extract_type == "triples":
             self.info_extract = InfoExtract(llm, example_prompt)
         elif extract_type == "property_graph":
-            self.property_graph_extract = PropertyGraphExtract(llm, example_prompt)
+            self.property_graph_extract = PropertyGraphExtract(
+                llm,
+                example_prompt,
+                max_workers=self.wk_input.graph_extract_max_workers,
+            )
         else:
             return CStatus(-1, f"Unsupported extract_type: {extract_type}")
         return super().node_init()

hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py

@@ -19,11 +19,13 @@
 
 import json
 import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Dict, List
 
 from hugegraph_llm.config import prompt
 from hugegraph_llm.document.chunk_split import ChunkSplitter
 from hugegraph_llm.models.llms.base import BaseLLM
+from hugegraph_llm.utils.graph_extract_config import validate_graph_extract_max_workers
 from hugegraph_llm.utils.log import log
 
 # TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable.
@@ -78,9 +80,15 @@ def filter_item(schema, items) -> List[Dict[str, Any]]:
 
 
 class PropertyGraphExtract:
-    def __init__(self, llm: BaseLLM, example_prompt: str = prompt.extract_graph_prompt) -> None:
+    def __init__(
+        self,
+        llm: BaseLLM,
+        example_prompt: str = prompt.extract_graph_prompt,
+        max_workers: int = 1,
+    ) -> None:
         self.llm = llm
         self.example_prompt = example_prompt
+        self.max_workers = validate_graph_extract_max_workers(max_workers)
         self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"}  # pylint: disable=invalid-name
 
     def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]:
@@ -90,26 +98,86 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]:
             context["vertices"] = []
         if "edges" not in context:
             context["edges"] = []
+
         items = []
-        for chunk in chunks:
-            proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk)
-            log.debug(
-                "[LLM] %s input: %s \n output:%s",
-                self.__class__.__name__,
-                chunk,
-                proceeded_chunk,
-            )
-            items.extend(self._extract_and_filter_label(schema, proceeded_chunk))
+        if self.max_workers == 1 or len(chunks) <= 1:
+            chunk_results = [
+                self._extract_chunk_items(schema, chunk, index, len(chunks)) for index, chunk in enumerate(chunks)
+            ]
+        else:
+            chunk_results = self._extract_chunks_concurrently(schema, chunks)
+
+        for chunk_items in chunk_results:
+            items.extend(chunk_items)
+
         items = filter_item(schema, items)
         for item in items:
             if item["type"] == "vertex":
                 context["vertices"].append(item)
             elif item["type"] == "edge":
                 context["edges"].append(item)
-
         context["call_count"] = context.get("call_count", 0) + len(chunks)
         return context
 
+    def _extract_chunks_concurrently(self, schema, chunks):
+        worker_count = min(self.max_workers, len(chunks))
+        chunk_results = [None] * len(chunks)
+        with ThreadPoolExecutor(max_workers=worker_count) as executor:
+            future_to_index = {
+                executor.submit(
+                    self._extract_chunk_items,
+                    schema,
+                    chunk,
+                    index,
+                    len(chunks),
+                ): index
+                for index, chunk in enumerate(chunks)
+            }
+            for future in as_completed(future_to_index):
+                index = future_to_index[future]
+                chunk_results[index] = future.result()
+        return chunk_results
+
+    def _extract_chunk_items(self, schema, chunk, chunk_index, chunk_count):
+        try:
+            proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk)
+            self._extract_property_graph_json(proceeded_chunk)
+            self._extract_property_graph_json(proceeded_chunk)
+        except Exception as exc:
+            raise RuntimeError(f"Graph extraction failed for chunk {chunk_index + 1}/{chunk_count}: {exc}") from exc
+
+        log.debug(
+            "[LLM] %s input: %s \noutput:%s",
+            self.__class__.__name__,
+            chunk,
+            proceeded_chunk,
+        )
+        return self._extract_and_filter_label(schema, proceeded_chunk)
+
+    @staticmethod
+    def _extract_property_graph_json(text):
+        text = re.sub(r"```\w*\n?", "", text)
+        text = re.sub(r"```", "", text)
+        text = text.strip()
+
+        json_match = re.search(r"(\{.*\}|\[.*\])", text, re.DOTALL)
+        if not json_match:
+            raise ValueError("No JSON found in property graph extraction response")
+
+        json_str = json_match.group(1).strip()
+        try:
+            property_graph = json.loads(json_str)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Invalid property graph JSON: {exc.msg}") from exc
+
+        if isinstance(property_graph, list):
+            return property_graph
+
+        if not (isinstance(property_graph, dict) and "vertices" in property_graph and "edges" in property_graph):
+            raise ValueError("Invalid property graph format; expecting 'vertices' and 'edges'")
+
+        return property_graph
+
     def extract_property_graph_by_llm(self, schema, chunk):
         prompt = generate_extract_property_graph_prompt(chunk, schema)
         if self.example_prompt is not None:

hugegraph-llm/src/hugegraph_llm/state/ai_state.py

@@ -24,6 +24,7 @@ class WkFlowInput(GParam):
     texts: Optional[Union[str, List[str]]] = None  # texts input used by ChunkSplit Node
     language: Optional[str] = None  # language configuration used by ChunkSplit Node
     split_type: Optional[str] = None  # split type used by ChunkSplit Node
+    graph_extract_max_workers: int = 1  # bounded chunk extraction concurrency
     example_prompt: Optional[str] = None  # need by graph information extract
     schema: Optional[str] = None  # Schema information requeired by SchemaNode
     # Request-scoped HugeGraph connection; None falls back to global huge_settings.
@@ -86,6 +87,7 @@ def reset(self, _: CStatus) -> None:
         self.texts = None
         self.language = None
         self.split_type = None
+        self.graph_extract_max_workers = 1
         self.example_prompt = None
         self.schema = None
         self.graph_client_config = None

hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py

@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+MAX_GRAPH_EXTRACT_WORKERS = 8
+
+
+def validate_graph_extract_max_workers(value) -> int:
+    try:
+        workers = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}"
+        ) from exc
+
+    if workers < 1 or workers > MAX_GRAPH_EXTRACT_WORKERS:
+        raise ValueError(f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}")
+
+    return workers