imbajin commented on code in PR #370:
URL: https://github.com/apache/hugegraph-ai/pull/370#discussion_r3608281619
##########
hugegraph-llm/src/hugegraph_llm/operators/llm_op/property_graph_extract.py:
##########
@@ -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)
Review Comment:
⚠️ Item-level normalization failures still escape without the promised chunk
context. The `try` above ends before this call, so a valid top-level graph
whose vertex has a non-mapping `properties` value can fail later at
`properties.get(...)` with a bare `AttributeError`, making it impossible to
identify the offending chunk. Please include filtering and normalization in the
chunk-level exception wrapper and add a regression case for malformed item
fields.
##########
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
Review Comment:
⚠️ This concurrency proof depends on the 20 ms sleep in `CountingLLM`. On a
loaded runner, the second executor worker can start after the first call
finishes, causing a correct `max_workers=2` implementation to report
`max_active == 1`. Please coordinate the first two calls with a barrier or
started/release events so their overlap is deterministic while retaining the
upper-bound assertion.
##########
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
[email protected]("raw_value", [True, False, "1.0"])
+def test_graph_extract_rejects_raw_invalid_worker_values(raw_value):
Review Comment:
⚠️ The new REST tests only cover rejected raw values; none sends a valid
non-default worker count through a successful POST and inspects the scheduler
call. Removing or misrouting the forwarding in `GraphExtractService` would
therefore remain green. Please add a mocked-scheduler request with
`graph_extract_max_workers: 4` and assert that exact value reaches
`schedule_flow` (plus the default if useful).
##########
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()
Review Comment:
⚠️ Releasing the slow call from the failing worker makes the lifecycle
assertion racy: an implementation that returns immediately after observing the
failure can still pass if the slow thread reaches `slow_finished` before the
assertion runs. Please execute `extractor.run()` in a separate test thread,
assert it has not returned while the slow call remains blocked, then release
the slow call and verify the extraction returns only afterward.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
