From 3719fb39a16bbdc0ff725dd5dbdb3d1f836dd58a Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Fri, 10 Jul 2026 18:38:30 +0800 Subject: [PATCH 1/8] feat: add bounded graph extraction concurrency --- .../config/models/base_prompt_config.py | 2 + .../demo/rag_demo/vector_graph_block.py | 22 ++ .../src/hugegraph_llm/flows/graph_extract.py | 10 + .../nodes/llm_node/extract_info.py | 6 +- .../llm_op/property_graph_extract.py | 63 +++++- .../src/hugegraph_llm/state/ai_state.py | 2 + .../hugegraph_llm/utils/graph_index_utils.py | 13 +- .../test_graph_extract_concurrency.py | 203 ++++++++++++++++++ 8 files changed, 308 insertions(+), 13 deletions(-) create mode 100644 hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index 7a7cfa7f9..941f5f0dd 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/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(), diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 9e5cf0b52..86ef57e82 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/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 ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index 4c96434a6..b780fe27c 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/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, ) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py index 9ac26970c..0ef65edd8 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py +++ b/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() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index 3e3974746..d3574630d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/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: diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 739588c56..665fb52ef 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/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 diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 78e9030d4..d363f2fbe 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/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) diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py new file mode 100644 index 000000000..7ba34fc38 --- /dev/null +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -0,0 +1,203 @@ +# 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. + +import json +import threading +import time + +import gradio as gr +import pytest + +from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract +from hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils import graph_index_utils + +SCHEMA = { + "vertexlabels": [ + { + "id": 1, + "name": "person", + "id_strategy": "PRIMARY_KEY", + "primary_keys": ["name"], + "nullable_keys": [], + "properties": ["name"], + } + ], + "edgelabels": [], +} + + +class CountingLLM: + def __init__(self, delay=0.02, fail_on=None): + self.delay = delay + self.fail_on = fail_on + self.active = 0 + self.max_active = 0 + self.lock = threading.Lock() + self.calls = [] + + def generate(self, prompt): + chunk = self._chunk_from_prompt(prompt) + with self.lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + self.calls.append(chunk) + if chunk == self.fail_on: + raise RuntimeError("boom") + time.sleep(self.delay) + return json.dumps( + { + "vertices": [ + { + "label": "person", + "type": "vertex", + "properties": {"name": chunk}, + } + ], + "edges": [], + } + ) + finally: + with self.lock: + self.active -= 1 + + @staticmethod + def _chunk_from_prompt(prompt): + for marker in ( + "first", + "second", + "third", + "later", + "bad", + "ok", + "a", + "b", + "c", + "d", + ): + if marker in prompt: + return marker + raise AssertionError(f"Could not identify chunk in prompt: {prompt}") + + +def test_property_graph_extract_respects_configured_concurrency_limit(): + llm = CountingLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + result = extractor.run({"schema": SCHEMA, "chunks": ["a", "b", "c", "d"]}) + + assert llm.max_active <= 2 + assert llm.max_active > 1 + assert result["call_count"] == 4 + + +def test_property_graph_extract_serial_mode_keeps_one_active_call(): + llm = CountingLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=1) + + result = extractor.run({"schema": SCHEMA, "chunks": ["a", "b", "c"]}) + + assert llm.max_active == 1 + assert result["call_count"] == 3 + + +def test_property_graph_extract_preserves_chunk_merge_order_with_concurrency(): + llm = CountingLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=3) + + result = extractor.run({"schema": SCHEMA, "chunks": ["first", "second", "third"]}) + + assert [item["properties"]["name"] for item in result["vertices"]] == [ + "first", + "second", + "third", + ] + + +def test_property_graph_extract_failed_chunk_reports_chunk_context(): + llm = CountingLLM(fail_on="bad") + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + with pytest.raises(RuntimeError, match="chunk 2/3"): + extractor.run({"schema": SCHEMA, "chunks": ["ok", "bad", "later"]}) + + +def test_graph_extract_flow_prepare_stores_positive_concurrency(): + flow = GraphExtractFlow() + prepared_input = WkFlowInput() + + flow.prepare( + prepared_input, + "{}", + ["doc"], + "prompt", + "property_graph", + graph_extract_max_workers=3, + ) + + assert prepared_input.graph_extract_max_workers == 3 + + +def test_graph_extract_flow_prepare_rejects_invalid_concurrency(): + flow = GraphExtractFlow() + + with pytest.raises(ValueError, match="positive integer"): + flow.prepare( + WkFlowInput(), + "{}", + ["doc"], + "prompt", + "property_graph", + graph_extract_max_workers=0, + ) + + +def test_extract_graph_helper_forwards_concurrency(monkeypatch): + calls = {} + + class DummyScheduler: + def schedule_flow(self, flow_name, *args, **kwargs): + calls["flow_name"] = flow_name + calls["kwargs"] = kwargs + return "ok" + + monkeypatch.setattr( + graph_index_utils, + "read_documents", + lambda input_file, input_text: ["doc"], + ) + monkeypatch.setattr( + graph_index_utils.SchedulerSingleton, + "get_instance", + lambda: DummyScheduler(), + ) + + result = graph_index_utils.extract_graph([], "", "{}", "prompt", "document", 4) + + assert result == "ok" + assert calls["kwargs"]["graph_extract_max_workers"] == 4 + + +def test_extract_graph_helper_rejects_invalid_concurrency(monkeypatch): + monkeypatch.setattr( + graph_index_utils, + "read_documents", + lambda input_file, input_text: ["doc"], + ) + + with pytest.raises(gr.Error, match="positive integer"): + graph_index_utils.extract_graph([], "", "{}", "prompt", "document", 0) From fd572e398191f01324b406a2931df7162daeb982 Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Sat, 11 Jul 2026 15:31:40 +0800 Subject: [PATCH 2/8] fix: address graph extraction concurrency review comments --- .../hugegraph_llm/api/graph_extract_api.py | 1 + .../api/models/graph_extract_requests.py | 8 +++ .../src/hugegraph_llm/flows/graph_extract.py | 9 +-- .../llm_op/property_graph_extract.py | 29 +++++++- .../utils/graph_extract_config.py | 30 ++++++++ .../hugegraph_llm/utils/graph_index_utils.py | 9 ++- .../test_graph_extract_concurrency.py | 69 ++++++++++++++++++- 7 files changed, 139 insertions(+), 16 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py diff --git a/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py b/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py index 6884412d7..0ac382d43 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py +++ b/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) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py index d3e2654a4..218f3c70f 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py +++ b/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.") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index b780fe27c..645d1d6ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/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 @@ -54,13 +55,7 @@ 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.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 diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index d3574630d..c67933227 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -25,6 +25,7 @@ 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. @@ -87,7 +88,7 @@ def __init__( ) -> None: self.llm = llm self.example_prompt = example_prompt - self.max_workers = max(1, int(max_workers or 1)) + 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]]: @@ -140,6 +141,8 @@ def _extract_chunks_concurrently(self, schema, chunks): 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 @@ -151,6 +154,30 @@ def _extract_chunk_items(self, schema, chunk, chunk_index, chunk_count): ) 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: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py new file mode 100644 index 000000000..fe1a92b13 --- /dev/null +++ b/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 diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index d363f2fbe..f027feb1c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -28,6 +28,7 @@ SPLIT_TYPE_DOCUMENT, VALID_SPLIT_TYPES, ) +from hugegraph_llm.utils.graph_extract_config import validate_graph_extract_max_workers from ..config import huge_settings from .hugegraph_utils import clean_hg_data @@ -96,11 +97,9 @@ 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") + graph_extract_max_workers = validate_graph_extract_max_workers(graph_extract_max_workers) + except ValueError as exc: + raise gr.Error(str(exc)) from exc try: schedule_kwargs = {"split_type": split_type} if graph_extract_max_workers != 1: diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index 7ba34fc38..a3f6c5cff 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -19,6 +19,7 @@ import gradio as gr import pytest +from pydantic import ValidationError from hugegraph_llm.flows.graph_extract import GraphExtractFlow from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract @@ -41,9 +42,10 @@ class CountingLLM: - def __init__(self, delay=0.02, fail_on=None): + def __init__(self, delay=0.02, fail_on=None, malformed_on=None): self.delay = delay self.fail_on = fail_on + self.malformed_on = malformed_on self.active = 0 self.max_active = 0 self.lock = threading.Lock() @@ -58,6 +60,8 @@ def generate(self, prompt): self.calls.append(chunk) if chunk == self.fail_on: raise RuntimeError("boom") + if chunk == self.malformed_on: + return "this is not json" time.sleep(self.delay) return json.dumps( { @@ -155,7 +159,7 @@ def test_graph_extract_flow_prepare_stores_positive_concurrency(): def test_graph_extract_flow_prepare_rejects_invalid_concurrency(): flow = GraphExtractFlow() - with pytest.raises(ValueError, match="positive integer"): + with pytest.raises(ValueError, match="between 1 and 8"): flow.prepare( WkFlowInput(), "{}", @@ -199,5 +203,64 @@ def test_extract_graph_helper_rejects_invalid_concurrency(monkeypatch): lambda input_file, input_text: ["doc"], ) - with pytest.raises(gr.Error, match="positive integer"): + with pytest.raises(gr.Error, match="between 1 and 8"): graph_index_utils.extract_graph([], "", "{}", "prompt", "document", 0) + + +def test_property_graph_extract_malformed_chunk_reports_chunk_context(): + llm = CountingLLM(malformed_on="bad") + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + with pytest.raises(RuntimeError, match="chunk 2/3"): + extractor.run({"schema": SCHEMA, "chunks": ["ok", "bad", "later"]}) + + +def test_graph_extract_flow_prepare_rejects_concurrency_above_backend_cap(): + flow = GraphExtractFlow() + + with pytest.raises(ValueError, match="between 1 and 8"): + flow.prepare( + WkFlowInput(), + "{}", + ["doc"], + "prompt", + "property_graph", + graph_extract_max_workers=9, + ) + + +def test_extract_graph_helper_rejects_concurrency_above_backend_cap(monkeypatch): + monkeypatch.setattr( + graph_index_utils, + "read_documents", + lambda input_file, input_text: ["doc"], + ) + + with pytest.raises(gr.Error, match="between 1 and 8"): + graph_index_utils.extract_graph( + [], + "", + "{}", + "prompt", + "document", + 9, + ) + + +def test_rest_graph_extract_request_accepts_and_validates_concurrency(): + from hugegraph_llm.api.models.graph_extract_requests import GraphExtractRequest + + request = GraphExtractRequest( + texts="doc", + schema=SCHEMA, + graph_extract_max_workers=4, + ) + + assert request.graph_extract_max_workers == 4 + + with pytest.raises(ValidationError): + GraphExtractRequest( + texts="doc", + schema=SCHEMA, + graph_extract_max_workers=9, + ) From 182e1a3c428a424cff2c64d5133838ba4fb1a8a2 Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Sun, 12 Jul 2026 15:31:43 +0800 Subject: [PATCH 3/8] fix: address graph extraction concurrency review comments --- .../llm_op/property_graph_extract.py | 131 +++++++----------- .../utils/graph_extract_config.py | 20 ++- .../test_graph_extract_concurrency.py | 72 ++++++++++ 3 files changed, 140 insertions(+), 83 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index c67933227..4e4675ccf 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -122,27 +122,54 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: 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( + executor = ThreadPoolExecutor(max_workers=worker_count) + future_to_index = {} + next_index = 0 + + try: + while next_index < len(chunks) and len(future_to_index) < worker_count: + future = executor.submit( self._extract_chunk_items, schema, - chunk, - index, + chunks[next_index], + next_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 + ) + future_to_index[future] = next_index + next_index += 1 + + while future_to_index: + for future in as_completed(future_to_index): + index = future_to_index.pop(future) + try: + chunk_results[index] = future.result() + except Exception: + for pending_future in future_to_index: + pending_future.cancel() + executor.shutdown(wait=False, cancel_futures=True) + raise + + if next_index < len(chunks): + next_future = executor.submit( + self._extract_chunk_items, + schema, + chunks[next_index], + next_index, + len(chunks), + ) + future_to_index[next_future] = next_index + next_index += 1 + break + except Exception: + raise + else: + executor.shutdown(wait=True) + 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) + property_graph = 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 @@ -152,7 +179,7 @@ def _extract_chunk_items(self, schema, chunk, chunk_index, chunk_count): chunk, proceeded_chunk, ) - return self._extract_and_filter_label(schema, proceeded_chunk) + return self._extract_and_filter_label(schema, property_graph) @staticmethod def _extract_property_graph_json(text): @@ -272,68 +299,16 @@ def _normalize_edges(self, edges, edge_label_map, vertex_label_map, vertex_id_ma normalized_edges.append(edge) return normalized_edges - def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: - # Strip markdown code blocks (e.g. ```json ... ```) - text = re.sub(r"```\w*\n?", "", text) - text = re.sub(r"```", "", text) - text = text.strip() + def _extract_and_filter_label(self, schema, property_graph): + _ = schema + if isinstance(property_graph, str): + property_graph = self._extract_property_graph_json(property_graph) - # Try to extract JSON (object or array) - json_match = re.search(r"(\{.*\}|\[.*\])", text, re.DOTALL) - if not json_match: - log.critical("Invalid property graph! No JSON found, please check the output format example in prompt.") - return [] - json_str = json_match.group(1).strip() + if isinstance(property_graph, list): + graph_items = property_graph + elif isinstance(property_graph, dict): + graph_items = property_graph["vertices"] + property_graph["edges"] + else: + raise ValueError("Invalid property graph format") - items = [] - try: - property_graph = json.loads(json_str) - # Handle flat array format: convert to {"vertices": [...], "edges": [...]} - if isinstance(property_graph, list): - vertices = [item for item in property_graph if isinstance(item, dict) and item.get("type") == "vertex"] - edges = [item for item in property_graph if isinstance(item, dict) and item.get("type") == "edge"] - property_graph = {"vertices": vertices, "edges": edges} - # Expect property_graph to be a dict with keys "vertices" and "edges" - if not (isinstance(property_graph, dict) and "vertices" in property_graph and "edges" in property_graph): - log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") - return items - - # Create sets for valid vertex and edge labels based on the schema - vertex_label_map = {vertex["name"]: vertex for vertex in schema["vertexlabels"]} - edge_label_map = {edge["name"]: edge for edge in schema["edgelabels"]} - vertex_label_set = set(vertex_label_map) - edge_label_set = set(edge_label_map) - - def process_items(item_list, valid_labels, item_type): - parsed_items = [] - for item in item_list: - if not isinstance(item, dict): - log.warning("Invalid property graph item type '%s'.", type(item)) - continue - item = dict(item) - item_type_value = item.get("type", item_type) - item["type"] = item_type_value - if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): - log.warning("Invalid item keys '%s'.", item.keys()) - continue - if item_type_value != item_type: - log.warning("Invalid %s type '%s' has been ignored.", item_type, item_type_value) - continue - if item["label"] not in valid_labels: - log.warning( - "Invalid %s label '%s' has been ignored.", - item_type, - item["label"], - ) - continue - parsed_items.append(item) - return parsed_items - - vertex_items = process_items(property_graph["vertices"], vertex_label_set, "vertex") - vertices, vertex_id_map = self._normalize_vertices(vertex_items, vertex_label_map) - edge_items = process_items(property_graph["edges"], edge_label_set, "edge") - edges = self._normalize_edges(edge_items, edge_label_map, vertex_label_map, vertex_id_map) - items = vertices + edges - except json.JSONDecodeError: - log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") - return items + return [item for item in graph_items if isinstance(item, dict) and self.NECESSARY_ITEM_KEYS.issubset(item)] diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py index fe1a92b13..3be1093fc 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py @@ -17,12 +17,22 @@ def validate_graph_extract_max_workers(value) -> int: - try: + if isinstance(value, bool): + raise ValueError(f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}") + + if isinstance(value, int): + workers = value + elif isinstance(value, float): + if not value.is_integer(): + raise ValueError(f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}") 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 + elif isinstance(value, str): + stripped_value = value.strip() + if not stripped_value or not stripped_value.lstrip("+-").isdigit(): + raise ValueError(f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}") + workers = int(stripped_value) + else: + raise ValueError(f"graph_extract_max_workers must be an integer between 1 and {MAX_GRAPH_EXTRACT_WORKERS}") 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}") diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index a3f6c5cff..396d657c2 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -25,6 +25,7 @@ from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract from hugegraph_llm.state.ai_state import WkFlowInput from hugegraph_llm.utils import graph_index_utils +from hugegraph_llm.utils.graph_extract_config import validate_graph_extract_max_workers SCHEMA = { "vertexlabels": [ @@ -264,3 +265,74 @@ def test_rest_graph_extract_request_accepts_and_validates_concurrency(): schema=SCHEMA, graph_extract_max_workers=9, ) + + +class BlockingLLM: + def __init__(self): + self.started = [] + self.release = threading.Event() + self.lock = threading.Lock() + + def generate(self, prompt): + chunk = CountingLLM._chunk_from_prompt(prompt) + with self.lock: + self.started.append(chunk) + + if chunk == "bad": + raise RuntimeError("boom") + if chunk == "slow": + self.release.wait(timeout=1) + + return json.dumps( + { + "vertices": [ + { + "label": "person", + "type": "vertex", + "properties": {"name": chunk}, + } + ], + "edges": [], + } + ) + + +def test_property_graph_extract_cancels_queued_chunks_after_failure(): + llm = BlockingLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + with pytest.raises(RuntimeError, match="chunk 2/4"): + extractor.run({"schema": SCHEMA, "chunks": ["slow", "bad", "later", "d"]}) + + llm.release.set() + time.sleep(0.05) + + assert "later" not in llm.started + assert "d" not in llm.started + + +def test_property_graph_extract_parses_successful_chunk_response_once(monkeypatch): + llm = CountingLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=1) + parse_count = 0 + original_parser = extractor._extract_property_graph_json + + def counting_parser(text): + nonlocal parse_count + parse_count += 1 + return original_parser(text) + + monkeypatch.setattr(extractor, "_extract_property_graph_json", counting_parser) + + extractor.run({"schema": SCHEMA, "chunks": ["a", "b"]}) + + assert parse_count == 2 + + +def test_graph_extract_worker_validator_rejects_fractional_and_bool_values(): + assert validate_graph_extract_max_workers(1.0) == 1 + assert validate_graph_extract_max_workers("8") == 8 + + for invalid_value in (True, False, 1.9, 8.9, "1.9"): + with pytest.raises(ValueError, match="between 1 and 8"): + validate_graph_extract_max_workers(invalid_value) From 5e3ac82b834dcc2d26e09fda50444f72c8805978 Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Sun, 12 Jul 2026 15:42:02 +0800 Subject: [PATCH 4/8] fix: restore graph extraction schema filtering --- .../llm_op/property_graph_extract.py | 52 ++++++++++++++++--- .../test_graph_extract_concurrency.py | 47 +++++++++++------ 2 files changed, 77 insertions(+), 22 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index 4e4675ccf..5cb2d46e4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -300,15 +300,53 @@ def _normalize_edges(self, edges, edge_label_map, vertex_label_map, vertex_id_ma return normalized_edges def _extract_and_filter_label(self, schema, property_graph): - _ = schema if isinstance(property_graph, str): property_graph = self._extract_property_graph_json(property_graph) if isinstance(property_graph, list): - graph_items = property_graph - elif isinstance(property_graph, dict): - graph_items = property_graph["vertices"] + property_graph["edges"] - else: - raise ValueError("Invalid property graph format") + property_graph = { + "vertices": [ + item for item in property_graph if isinstance(item, dict) and item.get("type") == "vertex" + ], + "edges": [item for item in property_graph if isinstance(item, dict) and item.get("type") == "edge"], + } + + 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 [item for item in graph_items if isinstance(item, dict) and self.NECESSARY_ITEM_KEYS.issubset(item)] + vertex_label_map = {vertex["name"]: vertex for vertex in schema["vertexlabels"]} + edge_label_map = {edge["name"]: edge for edge in schema["edgelabels"]} + vertex_label_set = set(vertex_label_map) + edge_label_set = set(edge_label_map) + + def process_items(item_list, valid_labels, item_type): + parsed_items = [] + for item in item_list: + if not isinstance(item, dict): + log.warning("Invalid property graph item type '%s'.", type(item)) + continue + + item = dict(item) + item_type_value = item.get("type", item_type) + item["type"] = item_type_value + + if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): + log.warning("Invalid item keys '%s'.", item.keys()) + continue + if item_type_value != item_type: + log.warning("Invalid %s type '%s' has been ignored.", item_type, item_type_value) + continue + if item["label"] not in valid_labels: + log.warning("Invalid %s label '%s' has been ignored.", item_type, item["label"]) + continue + + parsed_items.append(item) + return parsed_items + + vertex_items = process_items(property_graph["vertices"], vertex_label_set, "vertex") + vertices, vertex_id_map = self._normalize_vertices(vertex_items, vertex_label_map) + + edge_items = process_items(property_graph["edges"], edge_label_set, "edge") + edges = self._normalize_edges(edge_items, edge_label_map, vertex_label_map, vertex_id_map) + + return vertices + edges diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index 396d657c2..af5df5e63 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -14,6 +14,7 @@ # limitations under the License. import json +import re import threading import time @@ -82,21 +83,10 @@ def generate(self, prompt): @staticmethod def _chunk_from_prompt(prompt): - for marker in ( - "first", - "second", - "third", - "later", - "bad", - "ok", - "a", - "b", - "c", - "d", - ): - if marker in prompt: - return marker - raise AssertionError(f"Could not identify chunk in prompt: {prompt}") + match = re.search(r"## Text:\s*(.*?)\s*## Graph schema", prompt, re.DOTALL) + if not match: + raise AssertionError(f"Could not identify chunk in prompt: {prompt}") + return match.group(1).strip() def test_property_graph_extract_respects_configured_concurrency_limit(): @@ -336,3 +326,30 @@ def test_graph_extract_worker_validator_rejects_fractional_and_bool_values(): for invalid_value in (True, False, 1.9, 8.9, "1.9"): with pytest.raises(ValueError, match="between 1 and 8"): validate_graph_extract_max_workers(invalid_value) + + +def test_property_graph_extract_accepts_vertices_without_explicit_type(): + class UntypedVertexLLM: + def generate(self, prompt): + return json.dumps( + { + "vertices": [ + { + "label": "person", + "properties": {"name": "Ada"}, + }, + { + "label": "person", + "properties": {"name": "Bob"}, + }, + ], + "edges": [], + } + ) + + extractor = PropertyGraphExtract(UntypedVertexLLM(), example_prompt="", max_workers=1) + + result = extractor.run({"schema": SCHEMA, "chunks": ["a"]}) + + assert [vertex["label"] for vertex in result["vertices"]] == ["person", "person"] + assert [vertex["type"] for vertex in result["vertices"]] == ["vertex", "vertex"] From fea7c354bcf027362a4db06fd04555ccf9b394b0 Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Sun, 12 Jul 2026 15:46:35 +0800 Subject: [PATCH 5/8] fix: preserve graph extraction filter compatibility --- .../operators/llm_op/property_graph_extract.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index 5cb2d46e4..08dd7c96d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -301,7 +301,11 @@ def _normalize_edges(self, edges, edge_label_map, vertex_label_map, vertex_id_ma def _extract_and_filter_label(self, schema, property_graph): if isinstance(property_graph, str): - property_graph = self._extract_property_graph_json(property_graph) + try: + property_graph = self._extract_property_graph_json(property_graph) + except ValueError as exc: + log.warning("Invalid property graph extraction response: %s", exc) + return [] if isinstance(property_graph, list): property_graph = { From e80f434c65ec033c599aada9abd31f6c7997b74e Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Tue, 14 Jul 2026 14:17:31 +0800 Subject: [PATCH 6/8] fix: reject malformed graph extraction containers --- .../operators/llm_op/property_graph_extract.py | 6 ++++-- .../document/test_graph_extract_concurrency.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index 08dd7c96d..32f25396b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -181,8 +181,7 @@ def _extract_chunk_items(self, schema, chunk, chunk_index, chunk_count): ) return self._extract_and_filter_label(schema, property_graph) - @staticmethod - def _extract_property_graph_json(text): + def _extract_property_graph_json(self, text): text = re.sub(r"```\w*\n?", "", text) text = re.sub(r"```", "", text) text = text.strip() @@ -203,6 +202,9 @@ def _extract_property_graph_json(text): 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'") + if not isinstance(property_graph["vertices"], list) or not isinstance(property_graph["edges"], list): + raise ValueError("Invalid property graph format; 'vertices' and 'edges' must be lists") + return property_graph def extract_property_graph_by_llm(self, schema, chunk): diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index af5df5e63..7413f336b 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -353,3 +353,19 @@ def generate(self, prompt): assert [vertex["label"] for vertex in result["vertices"]] == ["person", "person"] assert [vertex["type"] for vertex in result["vertices"]] == ["vertex", "vertex"] + + +def test_property_graph_extract_malformed_container_shape_reports_chunk_context(): + class MalformedContainerLLM: + def generate(self, prompt): + return json.dumps( + { + "vertices": {}, + "edges": [], + } + ) + + extractor = PropertyGraphExtract(MalformedContainerLLM(), example_prompt="", max_workers=1) + + with pytest.raises(RuntimeError, match="chunk 1/1"): + extractor.run({"schema": SCHEMA, "chunks": ["a"]}) From 43d9bd6fa14369461c26a43845160ff13e7584d0 Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Wed, 15 Jul 2026 16:08:53 +0800 Subject: [PATCH 7/8] fix: tighten graph extraction worker validation --- .../api/models/graph_extract_requests.py | 11 +++- .../src/tests/api/test_graph_extract_api.py | 24 +++++++++ .../test_graph_extract_concurrency.py | 50 ++++++++++++++++--- 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py index 218f3c70f..978da30ff 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py @@ -21,7 +21,10 @@ 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 +from hugegraph_llm.utils.graph_extract_config import ( + MAX_GRAPH_EXTRACT_WORKERS, + validate_graph_extract_max_workers, +) class GraphExtractClientConfig(BaseModel): @@ -52,6 +55,12 @@ class GraphExtractRequest(BaseModel): 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.") diff --git a/hugegraph-llm/src/tests/api/test_graph_extract_api.py b/hugegraph-llm/src/tests/api/test_graph_extract_api.py index 2f7cdc3f5..318a9bfbb 100644 --- a/hugegraph-llm/src/tests/api/test_graph_extract_api.py +++ b/hugegraph-llm/src/tests/api/test_graph_extract_api.py @@ -256,6 +256,30 @@ def test_service_extract_sync_maps_errors_to_500(mock_singleton): assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR +@pytest.mark.parametrize("raw_value", [True, False, "1.0"]) +def test_graph_extract_rejects_raw_invalid_worker_values(raw_value): + response = _graph_client().post( + "/graph/extract", + json={ + "texts": "x", + "schema": INLINE_SCHEMA, + "graph_extract_max_workers": raw_value, + }, + ) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + + +@pytest.mark.parametrize("raw_value", [True, False, "1.0"]) +def test_request_model_rejects_raw_invalid_worker_values(raw_value): + with pytest.raises(ValidationError): + GraphExtractRequest( + texts="hello", + schema=INLINE_SCHEMA, + graph_extract_max_workers=raw_value, + ) + + def test_request_model_validation(): req = GraphExtractRequest(texts="hello", schema=INLINE_SCHEMA) assert req.texts == ["hello"] diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index 7413f336b..15992e39a 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -111,15 +111,51 @@ def test_property_graph_extract_serial_mode_keeps_one_active_call(): def test_property_graph_extract_preserves_chunk_merge_order_with_concurrency(): - llm = CountingLLM() - extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=3) + class ReverseFinishLLM: + def __init__(self): + self.first_started = threading.Event() + self.release_first = threading.Event() + self.completion_order = [] + self.lock = threading.Lock() + + def generate(self, prompt): + if "FIRST_CHUNK" in prompt: + chunk_name = "FIRST_CHUNK" + self.first_started.set() + self.release_first.wait(timeout=2) + time.sleep(0.05) + elif "SECOND_CHUNK" in prompt: + chunk_name = "SECOND_CHUNK" + assert self.first_started.wait(timeout=2) + self.release_first.set() + else: + raise AssertionError(f"Unexpected prompt: {prompt}") + + with self.lock: + self.completion_order.append(chunk_name) + + return json.dumps( + { + "vertices": [ + { + "label": "person", + "type": "vertex", + "properties": {"name": chunk_name}, + } + ], + "edges": [], + } + ) + + llm = ReverseFinishLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) - result = extractor.run({"schema": SCHEMA, "chunks": ["first", "second", "third"]}) + result = extractor.run({"schema": SCHEMA, "chunks": ["FIRST_CHUNK", "SECOND_CHUNK"]}) - assert [item["properties"]["name"] for item in result["vertices"]] == [ - "first", - "second", - "third", + assert llm.completion_order == ["SECOND_CHUNK", "FIRST_CHUNK"] + assert [vertex["properties"]["name"] for vertex in result["vertices"]] == [ + "FIRST_CHUNK", + "SECOND_CHUNK", ] From 8f43b448a1da162f3788d1669f623c2a7e8432bf Mon Sep 17 00:00:00 2001 From: nannan-2026 <1794949109@qq.com> Date: Fri, 17 Jul 2026 11:24:41 +0800 Subject: [PATCH 8/8] fix: handle graph extraction concurrency failures deterministically --- .../llm_op/property_graph_extract.py | 117 ++++++++++++------ .../test_graph_extract_concurrency.py | 85 +++++++++++++ 2 files changed, 162 insertions(+), 40 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py index 32f25396b..593270ee7 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py @@ -19,7 +19,7 @@ import json import re -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from typing import Any, Dict, List from hugegraph_llm.config import prompt @@ -122,49 +122,86 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: def _extract_chunks_concurrently(self, schema, chunks): worker_count = min(self.max_workers, len(chunks)) chunk_results = [None] * len(chunks) - executor = ThreadPoolExecutor(max_workers=worker_count) future_to_index = {} next_index = 0 - - try: + first_error = None + + def submit_next(executor): + nonlocal next_index + future = executor.submit( + self._extract_chunk_items, + schema, + chunks[next_index], + next_index, + len(chunks), + ) + future_to_index[future] = next_index + next_index += 1 + + def collect_done(done_futures): + completed_results = [] + completed_errors = [] + + for future in done_futures: + index = future_to_index.pop(future) + try: + completed_results.append((index, future.result())) + except Exception as exc: + completed_errors.append((index, exc)) + + return completed_results, completed_errors + + def update_first_error(completed_errors): + nonlocal first_error + if not completed_errors: + return + current_error = min(completed_errors, key=lambda item: item[0]) + if first_error is None or current_error[0] < first_error[0]: + first_error = current_error + + with ThreadPoolExecutor(max_workers=worker_count) as executor: while next_index < len(chunks) and len(future_to_index) < worker_count: - future = executor.submit( - self._extract_chunk_items, - schema, - chunks[next_index], - next_index, - len(chunks), - ) - future_to_index[future] = next_index - next_index += 1 - - while future_to_index: - for future in as_completed(future_to_index): - index = future_to_index.pop(future) - try: - chunk_results[index] = future.result() - except Exception: - for pending_future in future_to_index: - pending_future.cancel() - executor.shutdown(wait=False, cancel_futures=True) - raise - - if next_index < len(chunks): - next_future = executor.submit( - self._extract_chunk_items, - schema, - chunks[next_index], - next_index, - len(chunks), - ) - future_to_index[next_future] = next_index - next_index += 1 + submit_next(executor) + + while future_to_index and first_error is None: + done, _ = wait(future_to_index, return_when=FIRST_COMPLETED) + # Drain any other futures that completed before we decide whether to + # submit more work. This prevents submitting queued chunks after an + # already-visible in-flight failure. + done.update(future for future in future_to_index if future.done()) + + completed_results, completed_errors = collect_done(done) + update_first_error(completed_errors) + + for index, result in completed_results: + chunk_results[index] = result + + if first_error is not None: break - except Exception: - raise - else: - executor.shutdown(wait=True) - return chunk_results + + while next_index < len(chunks) and len(future_to_index) < worker_count: + submit_next(executor) + + if first_error is not None: + # Do not submit more chunks after the first observed failure. + # Cancel queued work if any, then wait for already-running calls so no + # background LLM call outlives this failed extraction. + for future in future_to_index: + future.cancel() + + remaining_errors = [] + for future, index in list(future_to_index.items()): + if future.cancelled(): + continue + try: + future.result() + except Exception as exc: + remaining_errors.append((index, exc)) + + update_first_error(remaining_errors) + raise first_error[1] + + return chunk_results def _extract_chunk_items(self, schema, chunk, chunk_index, chunk_count): try: diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py index 15992e39a..73d51abc2 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -405,3 +405,88 @@ def generate(self, prompt): with pytest.raises(RuntimeError, match="chunk 1/1"): extractor.run({"schema": SCHEMA, "chunks": ["a"]}) + + +def test_property_graph_extract_waits_for_in_flight_call_after_failure(): + class LifecycleLLM: + def __init__(self): + self.slow_started = threading.Event() + self.release_slow = threading.Event() + self.slow_finished = threading.Event() + self.calls = [] + self.lock = threading.Lock() + + def generate(self, prompt): + with self.lock: + if "LATER_CHUNK" in prompt: + self.calls.append("LATER_CHUNK") + elif "SLOW_CHUNK" in prompt: + self.calls.append("SLOW_CHUNK") + elif "FAIL_CHUNK" in prompt: + self.calls.append("FAIL_CHUNK") + + if "SLOW_CHUNK" in prompt: + self.slow_started.set() + self.release_slow.wait(timeout=2) + self.slow_finished.set() + return json.dumps({"vertices": [], "edges": []}) + + if "FAIL_CHUNK" in prompt: + assert self.slow_started.wait(timeout=2) + self.release_slow.set() + raise ValueError("fast failure") + + if "LATER_CHUNK" in prompt: + return json.dumps({"vertices": [], "edges": []}) + + raise AssertionError(f"Unexpected prompt: {prompt}") + + llm = LifecycleLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + with pytest.raises(RuntimeError, match="chunk 2/3"): + extractor.run({"schema": SCHEMA, "chunks": ["SLOW_CHUNK", "FAIL_CHUNK", "LATER_CHUNK"]}) + + assert llm.slow_finished.is_set() + assert "LATER_CHUNK" not in llm.calls + + +def test_property_graph_extract_reports_lowest_in_flight_failure_index(): + class FailureOrderLLM: + def __init__(self): + self.first_started = threading.Event() + self.second_failed = threading.Event() + self.calls = [] + self.lock = threading.Lock() + + def generate(self, prompt): + with self.lock: + if "LATER_CHUNK" in prompt: + self.calls.append("LATER_CHUNK") + elif "FIRST_FAIL_CHUNK" in prompt: + self.calls.append("FIRST_FAIL_CHUNK") + elif "SECOND_FAIL_CHUNK" in prompt: + self.calls.append("SECOND_FAIL_CHUNK") + + if "FIRST_FAIL_CHUNK" in prompt: + self.first_started.set() + assert self.second_failed.wait(timeout=2) + raise ValueError("first failure") + + if "SECOND_FAIL_CHUNK" in prompt: + assert self.first_started.wait(timeout=2) + self.second_failed.set() + raise ValueError("second failure") + + if "LATER_CHUNK" in prompt: + return json.dumps({"vertices": [], "edges": []}) + + raise AssertionError(f"Unexpected prompt: {prompt}") + + llm = FailureOrderLLM() + extractor = PropertyGraphExtract(llm, example_prompt="", max_workers=2) + + with pytest.raises(RuntimeError, match="chunk 1/3"): + extractor.run({"schema": SCHEMA, "chunks": ["FIRST_FAIL_CHUNK", "SECOND_FAIL_CHUNK", "LATER_CHUNK"]}) + + assert "LATER_CHUNK" not in llm.calls