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,11 @@
 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,
+    validate_graph_extract_max_workers,
+)
+
 
 class GraphExtractClientConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -44,6 +49,18 @@ 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.",
+    )
+
+    @field_validator("graph_extract_max_workers", mode="before")
+    @classmethod
+    def validate_graph_extract_max_workers_field(cls, value):
+        return validate_graph_extract_max_workers(value)
+
     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()