Skip to content

feat: add bounded graph extraction concurrency #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nw9663644-eng wants to merge 8 commits into
base: main
Choose a base branch
from
Open

feat: add bounded graph extraction concurrency #370

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3719fb3
feat: add bounded graph extraction concurrency
wn12222 Jul 10, 2026
fd572e3
fix: address graph extraction concurrency review comments
wn12222 Jul 11, 2026
182e1a3
fix: address graph extraction concurrency review comments
wn12222 Jul 12, 2026
5e3ac82
fix: restore graph extraction schema filtering
wn12222 Jul 12, 2026
fea7c35
fix: preserve graph extraction filter compatibility
wn12222 Jul 12, 2026
e80f434
fix: reject malformed graph extraction containers
wn12222 Jul 14, 2026
43d9bd6
fix: tighten graph extraction worker validation
wn12222 Jul 15, 2026
8f43b44
fix: handle graph extraction concurrency failures deterministically
wn12222 Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
8 changes: 8 additions & 0 deletions hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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(

@imbajin imbajin Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 REST validation coerces some raw inputs before the shared validator can enforce its contract. A direct probe shows graph_extract_max_workers=true and "1.0" are both accepted here as integer 1, while validate_graph_extract_max_workers() rejects those same raw values. This makes REST behavior inconsistent with the flow/helper entry points and can silently turn malformed input into a valid request. Please add a mode="before" field validator (or equivalent strict raw-input validation) that applies the shared rules, and cover boolean and decimal-string request values with regression tests.

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.")

Expand Down
2 changes: 2 additions & 0 deletions hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ""
Expand Down Expand Up @@ -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(),
Expand Down
22 changes: 22 additions & 0 deletions hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()


Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -494,6 +511,7 @@ def create_vector_graph_block():
input_schema,
info_extract_template,
graph_split_type,
graph_extract_max_workers,
],
)

Expand All @@ -506,6 +524,7 @@ def create_vector_graph_block():
input_schema,
info_extract_template,
graph_split_type,
graph_extract_max_workers,
],
outputs=[out],
).then(
Expand All @@ -515,6 +534,7 @@ def create_vector_graph_block():
input_schema,
info_extract_template,
graph_split_type,
graph_extract_max_workers,
],
)

Expand All @@ -527,6 +547,7 @@ def create_vector_graph_block():
input_schema,
info_extract_template,
graph_split_type,
graph_extract_max_workers,
],
)

Expand All @@ -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
)

Expand Down
5 changes: 5 additions & 0 deletions hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -44,6 +45,7 @@ def prepare(
extract_type,
split_type=SPLIT_TYPE_DOCUMENT,
language="zh",
graph_extract_max_workers=1,
**kwargs,
):
# prepare input data
Expand All @@ -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
Expand All @@ -76,6 +79,7 @@ def build_flow(
extract_type,
split_type=SPLIT_TYPE_DOCUMENT,
language="zh",
graph_extract_max_workers=1,
**kwargs,
):
pipeline = GPipeline()
Expand All @@ -89,6 +93,7 @@ def build_flow(
extract_type,
split_type,
language,
graph_extract_max_workers,
**kwargs,
)

Expand Down
6 changes: 5 additions & 1 deletion hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading
Loading