feat: add bounded graph extraction concurrency

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

@@ -44,6 +44,7 @@ def prepare(
         extract_type,
         split_type=SPLIT_TYPE_DOCUMENT,
         language="zh",
+        graph_extract_max_workers=1,
         **kwargs,
     ):
         # prepare input data
@@ -53,6 +54,13 @@ def prepare(
             raise ValueError("split_type must be document, paragraph, or sentence")
 
         prepared_input.split_type = split_type
+        try:
+            graph_extract_max_workers = int(graph_extract_max_workers)
+        except (TypeError, ValueError) as exc:
+            raise ValueError("graph_extract_max_workers must be a positive integer") from exc
+        if graph_extract_max_workers < 1:
+            raise ValueError("graph_extract_max_workers must be a positive integer")
+        prepared_input.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 +84,7 @@ def build_flow(
         extract_type,
         split_type=SPLIT_TYPE_DOCUMENT,
         language="zh",
+        graph_extract_max_workers=1,
         **kwargs,
     ):
         pipeline = GPipeline()
@@ -89,6 +98,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,6 +19,7 @@
 
 import json
 import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Dict, List
 
 from hugegraph_llm.config import prompt
@@ -78,9 +79,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 = max(1, int(max_workers or 1))
         self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"}  # pylint: disable=invalid-name
 
     def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]:
@@ -90,26 +97,60 @@ 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)
+        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)
+
     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_index_utils.py

@@ -87,6 +87,7 @@ def extract_graph(
     schema,
     example_prompt,
     split_type=SPLIT_TYPE_DOCUMENT,
+    graph_extract_max_workers=1,
 ) -> str:
     texts = read_documents(input_file, input_text)
     scheduler = SchedulerSingleton.get_instance()
@@ -95,13 +96,23 @@ def extract_graph(
     if split_type not in VALID_SPLIT_TYPES:
         raise gr.Error("split_type must be document, paragraph, or sentence")
     try:
+        graph_extract_max_workers = int(graph_extract_max_workers)
+    except (TypeError, ValueError) as exc:
+        raise gr.Error("graph_extract_max_workers must be a positive integer") from exc
+    if graph_extract_max_workers < 1:
+        raise gr.Error("graph_extract_max_workers must be a positive integer")
+    try:
+        schedule_kwargs = {"split_type": split_type}
+        if graph_extract_max_workers != 1:
+            schedule_kwargs["graph_extract_max_workers"] = graph_extract_max_workers
+
         return scheduler.schedule_flow(
             FlowName.GRAPH_EXTRACT,
             schema,
             texts,
             example_prompt,
             "property_graph",
-            split_type=split_type,
+            **schedule_kwargs,
         )
     except Exception as e:  # pylint: disable=broad-exception-caught
         log.error(e)