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..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,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.") 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..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 @@ -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, ) 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..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,11 +19,13 @@ import json import re +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait 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,152 @@ 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) + future_to_index = {} + next_index = 0 + 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: + 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 + + 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: + proceeded_chunk = self.extract_property_graph_by_llm(schema, 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 + + log.debug( + "[LLM] %s input: %s \noutput:%s", + self.__class__.__name__, + chunk, + proceeded_chunk, + ) + return self._extract_and_filter_label(schema, property_graph) + + def _extract_property_graph_json(self, 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'") + + 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): prompt = generate_extract_property_graph_prompt(chunk, schema) if self.example_prompt is not None: @@ -204,68 +338,58 @@ 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): + if isinstance(property_graph, str): + try: + property_graph = self._extract_property_graph_json(property_graph) + except ValueError as exc: + log.warning("Invalid property graph extraction response: %s", exc) + return [] - # 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): + 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"], + } - 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 + 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'") + + 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/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_extract_config.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py new file mode 100644 index 000000000..3be1093fc --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_extract_config.py @@ -0,0 +1,40 @@ +# 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: + 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) + 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}") + + 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 78e9030d4..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 @@ -87,6 +88,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 +97,21 @@ 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 = 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: + 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/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 new file mode 100644 index 000000000..73d51abc2 --- /dev/null +++ b/hugegraph-llm/src/tests/document/test_graph_extract_concurrency.py @@ -0,0 +1,492 @@ +# 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 re +import threading +import time + +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 +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": [ + { + "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, 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() + 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") + if chunk == self.malformed_on: + return "this is not json" + 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): + 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(): + 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(): + 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_CHUNK", "SECOND_CHUNK"]}) + + assert llm.completion_order == ["SECOND_CHUNK", "FIRST_CHUNK"] + assert [vertex["properties"]["name"] for vertex in result["vertices"]] == [ + "FIRST_CHUNK", + "SECOND_CHUNK", + ] + + +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="between 1 and 8"): + 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="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, + ) + + +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) + + +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"] + + +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"]}) + + +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