imbajin commented on code in PR #361:
URL: https://github.com/apache/hugegraph-ai/pull/361#discussion_r3492062164
##########
hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py:
##########
@@ -15,55 +15,352 @@
# specific language governing permissions and limitations
# under the License.
-import json
+from typing import Optional
-from fastapi import APIRouter, HTTPException, status
+from fastapi import APIRouter, HTTPException, Request, status
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from fastapi.routing import APIRoute
-from hugegraph_llm.api.models.graph_extract_requests import GraphExtractRequest
-from hugegraph_llm.api.models.graph_extract_responses import
GraphExtractResponse
-from hugegraph_llm.config import prompt
-from hugegraph_llm.flows import FlowName
-from hugegraph_llm.flows.scheduler import SchedulerSingleton
+from hugegraph_llm.api.models.graph_extract_requests import (
+ GraphExtractAndImportRequest,
+ GraphExtractRequest,
+ GraphImportRequest,
+)
+from hugegraph_llm.api.models.graph_extract_responses import (
+ GraphExtractAndImportResponse,
+ GraphExtractError,
+ GraphExtractJobCreateResponse,
+ GraphExtractJobStatusResponse,
+ GraphExtractResponse,
+ GraphImportResponse,
+)
+from hugegraph_llm.services.graph_extract_jobs import (
+ GraphExtractJob,
+ GraphExtractJobStatus,
+ InMemoryGraphExtractJobStore,
+)
+from hugegraph_llm.services.graph_extract_service import (
+ FlowOutputValidationError,
+ GraphExtractService,
+ GraphImportService,
+)
from hugegraph_llm.utils.log import log
+GRAPH_EXTRACT_FLOW_OUTPUT_ERROR = "Graph extraction flow output is invalid"
+GRAPH_EXTRACT_RUNTIME_ERROR = "Graph extraction failed during execution"
+GRAPH_IMPORT_FLOW_OUTPUT_ERROR = "Graph import flow output is invalid"
+GRAPH_IMPORT_RUNTIME_ERROR = "Graph import failed during execution"
-class GraphExtractService:
- @staticmethod
- def extract_sync(req: GraphExtractRequest) -> GraphExtractResponse:
+
+def _error(code: str, message: str, phase: str, job_id: Optional[str] = None)
-> dict:
+ return GraphExtractError(code=code, message=message, phase=phase,
job_id=job_id).model_dump(exclude_none=True)
+
+
+def _job_ts(value) -> Optional[str]:
+ return value.isoformat() if value else None
+
+
+def _job_status_response(job: GraphExtractJob) ->
GraphExtractJobStatusResponse:
+ return GraphExtractJobStatusResponse(
+ job_id=job.job_id,
+ status=job.status,
+ created_at=_job_ts(job.created_at),
+ updated_at=_job_ts(job.updated_at),
+ started_at=_job_ts(job.started_at),
+ finished_at=_job_ts(job.finished_at),
+ expires_at=_job_ts(job.expires_at),
+ error=job.error,
+ )
+
+
+def _validation_message(errors) -> str:
+ details = []
+ for error in errors:
+ loc = ".".join(str(part) for part in error.get("loc", []) if part not
in {"body"})
+ msg = error.get("msg", "invalid input")
+ err_type = error.get("type", "validation_error")
+ details.append(f"{loc or 'request'}: {msg} ({err_type})")
+ return "; ".join(details) or "request validation failed"
+
+
+class GraphExtractAPIRoute(APIRoute):
+ def get_route_handler(self):
+ original_route_handler = super().get_route_handler()
+
+ async def custom_route_handler(request: Request):
+ try:
+ return await original_route_handler(request)
+ except RequestValidationError as exc:
+ return JSONResponse(
+ status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+ content={
+ "detail": _error(
Review Comment:
⚠️ This route class is installed while registering extract, job, import, and
extract-and-import endpoints, but every request-validation failure is
hard-coded as `GRAPH_EXTRACT_VALIDATION_ERROR`. Invalid `/graph/import` bodies
therefore return an extract-specific code/phase, making client error handling
ambiguous. Please choose the validation error code/phase based on the route, or
use separate handlers for extract and import routes, and cover invalid import
payloads in endpoint tests.
##########
hugegraph-llm/src/hugegraph_llm/services/graph_extract_service.py:
##########
@@ -0,0 +1,376 @@
+# 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 time
+from typing import Any, Dict, List, Optional
+
+from hugegraph_llm.api.models.graph_extract_requests import (
+ GraphExtractAndImportRequest,
+ GraphExtractRequest,
+ GraphImportRequest,
+ SchemaInput,
+ _validate_schema_value,
+)
+from hugegraph_llm.api.models.graph_extract_responses import (
+ GraphExtractResponse,
+ GraphImportResponse,
+)
+from hugegraph_llm.config import prompt
+from hugegraph_llm.flows import FlowName
+from hugegraph_llm.flows.scheduler import SchedulerSingleton
+from hugegraph_llm.utils.log import log
+
+SENSITIVE_CLIENT_CONFIG_KEYS = {"pwd", "password", "token", "api_key",
"secret"}
+SAFE_IMPORT_ERROR_KEYS = {"kind", "index", "label", "key", "reason"}
+
+
+class FlowOutputValidationError(ValueError):
+ """Raised when a workflow returns malformed output."""
+
+
+def normalize_schema(schema: SchemaInput) -> str:
+ schema = _validate_schema_value(schema)
+ if isinstance(schema, dict):
+ return json.dumps(schema, ensure_ascii=False)
+
+ schema_text = str(schema).strip()
+ if schema_text.startswith("{"):
+ try:
+ parsed_schema = json.loads(schema_text)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"schema must be valid JSON: {exc.msg}") from exc
+ return json.dumps(parsed_schema, ensure_ascii=False)
+ return schema_text
+
+
+def _redact_client_config(client_config) -> Dict[str, Any]:
+ if client_config is None:
+ return {}
+ config = client_config if isinstance(client_config, dict) else
client_config.model_dump(exclude_none=True)
+ return {key: ("***" if key in SENSITIVE_CLIENT_CONFIG_KEYS and value else
value) for key, value in config.items()}
+
+
+def _schema_graph_name(schema: str) -> Optional[str]:
+ schema_text = str(schema).strip()
+ return None if schema_text.startswith("{") else schema_text
+
+
+def apply_client_config(
+ client_config,
+ schema: Optional[str] = None,
+ align_graph_with_schema: bool = False,
+) -> Optional[Dict[str, Any]]:
+ if client_config is None:
+ config = {}
+ elif isinstance(client_config, dict):
+ config = {key: value for key, value in client_config.items() if value
is not None}
+ else:
+ config = client_config.model_dump(exclude_none=True)
+ schema_graph = _schema_graph_name(schema) if schema else None
+ if schema_graph:
+ target_graph = config.get("graph")
+ if target_graph and target_graph != schema_graph:
+ raise ValueError("schema graph name must match
client_config.graph")
+ if align_graph_with_schema:
+ config["graph"] = schema_graph
+ return config or None
+
+
+def _parse_flow_json(raw_result: Any, error_message: str) -> Dict[str, Any]:
+ if isinstance(raw_result, dict):
+ return raw_result
+ if not isinstance(raw_result, str):
+ raise FlowOutputValidationError(error_message)
+ try:
+ parsed = json.loads(raw_result)
+ except json.JSONDecodeError as exc:
+ raise FlowOutputValidationError(error_message) from exc
+ if not isinstance(parsed, dict):
+ raise FlowOutputValidationError(error_message)
+ return parsed
+
+
+def _pop_warnings(result: Dict[str, Any]) -> List[str]:
+ warnings = []
+ warning = result.pop("warning", None)
+ if warning:
+ warnings.append(str(warning))
+ extra_warnings = result.pop("warnings", None)
+ if isinstance(extra_warnings, list):
+ warnings.extend(str(item) for item in extra_warnings)
+ elif extra_warnings:
+ warnings.append(str(extra_warnings))
+ return warnings
+
+
+def _count_items(result: Dict[str, Any], key: str) -> int:
+ value = result.get(key)
+ return len(value) if isinstance(value, list) else 0
+
+
+def _validate_property_graph_result(result: Dict[str, Any]) -> None:
+ vertices = result.get("vertices", [])
+ edges = result.get("edges", [])
+ if not isinstance(vertices, list) or not isinstance(edges, list):
+ raise FlowOutputValidationError("property graph result must contain
list vertices and edges")
+ for vertex in vertices:
+ if not isinstance(vertex, dict) or "label" not in vertex or
"properties" not in vertex:
Review Comment:
⚠️ The `/graph/extract` workflow-output contract is weaker than the import
request contract. This only checks that each vertex/edge is a dict and has the
required keys, but it does not verify that `label/outV/outVLabel/inV/inVLabel`
are non-empty strings or that `properties` is an object. A scheduler result
like `{"vertices":[{"label":"person","properties":null}],"edges":[]}` would be
returned as a successful extract response even though `/graph/import` rejects
the same property-graph payload. Please align this validator with
`GraphImportRequest.validate_data()` and add a contract test for malformed
workflow output.
##########
hugegraph-llm/src/tests/api/test_graph_extract_jobs.py:
##########
@@ -0,0 +1,243 @@
+# 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.
+
+from concurrent.futures import ThreadPoolExecutor
+from datetime import timedelta
+from queue import Full
+from unittest.mock import Mock
+
+from fastapi import APIRouter, FastAPI, status
+from fastapi.testclient import TestClient
+
+from hugegraph_llm.api.graph_extract_api import graph_extract_http_api
+from hugegraph_llm.api.models.graph_extract_responses import
GraphExtractResponse
+from hugegraph_llm.services.graph_extract_jobs import GraphExtractJobStatus,
InMemoryGraphExtractJobStore
+
+
+def _payload():
+ return {
+ "texts": ["marko knows vadas"],
+ "schema": {
+ "vertexlabels": [{"name": "person", "properties": ["name"]}],
+ "edgelabels": [{"name": "knows", "source_label": "person",
"target_label": "person"}],
+ },
+ "example_prompt": "extract graph",
+ }
+
+
+def _client(service=None, job_store=None, run_jobs_inline=True):
Review Comment:
⚠️ The production default for `graph_extract_http_api(...)` is
`run_jobs_inline=None`, which submits work through `jobs.submit_job()` and
daemon workers, but this helper defaults to inline execution and the tests only
exercise inline mode or a deliberately non-running pending mode. Please add a
route-level test that omits `run_jobs_inline`, creates a job, polls until it
reaches a terminal state, and verifies result retrieval so the public async
worker path is covered.
##########
hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph.py:
##########
@@ -303,6 +321,293 @@ def test_init_schema_if_need(self,
mock_handle_graph_creation, mock_create_prope
# Verify that edgeLabel was called for each edge label
self.assertEqual(schema_mocks["edge_label"].call_count, 1) # 1 edge
label
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._check_property_data_type")
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation")
+ def test_load_into_graph(self, mock_handle_graph_creation,
mock_check_property_data_type):
+ """Test load_into_graph method."""
+ # Setup mocks
+ mock_handle_graph_creation.return_value = MagicMock(id="vertex_id")
+ mock_check_property_data_type.return_value = True
+
+ # Create vertices with proper data types according to schema
+ vertices = [
+ {"label": "person", "properties": {"name": "Tom Hanks", "age":
67}},
+ {"label": "movie", "properties": {"title": "Forrest Gump", "year":
1994}},
+ ]
+
+ edges = [
+ {
+ "label": "acted_in",
+ "properties": {"role": "Forrest Gump"},
+ "outV": "person:Tom Hanks", # Use the format expected by the
implementation
+ "inV": "movie:Forrest Gump", # Use the format expected by the
implementation
+ }
+ ]
+
+ # Call the method
+ self.commit2graph.load_into_graph(vertices, edges, self.schema)
+
+ # Verify that _handle_graph_creation was called for each vertex and
edge
+ self.assertEqual(mock_handle_graph_creation.call_count, 3) # 2
vertices + 1 edge
+
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._check_property_data_type")
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation")
+ def test_load_into_graph_returns_actual_created_and_skipped_counts(
+ self, mock_handle_graph_creation, mock_check_property_data_type
+ ):
+ """Test import stats reflect actual created and skipped graph
elements."""
+ mock_handle_graph_creation.side_effect = [
+ MagicMock(id="person:Tom Hanks"),
+ None,
+ ]
+ mock_check_property_data_type.return_value = True
+ vertices = [
+ {"label": "person", "properties": {"name": "Tom Hanks", "age":
67}},
+ {"label": "unknown", "properties": {"name": "Ignored"}},
+ ]
+ edges = [
+ {
+ "label": "acted_in",
+ "properties": {"role": "Forrest Gump"},
+ "outV": "person:Tom Hanks",
+ "inV": "movie:Forrest Gump",
+ }
+ ]
+
+ result = self.commit2graph.load_into_graph(vertices, edges,
self.schema)
+
+ self.assertEqual(result["vertices_attempted"], 2)
+ self.assertEqual(result["vertices_created"], 1)
+ self.assertEqual(result["vertices_skipped"], 1)
+ self.assertEqual(result["edges_attempted"], 1)
+ self.assertEqual(result["edges_created"], 0)
+ self.assertEqual(result["edges_skipped"], 1)
+ self.assertIn(
+ {"kind": "vertex", "index": 1, "reason": "vertex_label_not_found",
"label": "unknown"},
+ result["errors"],
+ )
+
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation")
+ def test_load_into_graph_with_data_type_validation_success(self,
mock_handle_graph_creation):
+ """Test load_into_graph method with successful data type validation."""
+ # Setup mocks
+ mock_handle_graph_creation.return_value = MagicMock(id="vertex_id")
+
+ # Create vertices with correct data types matching schema expectations
+ vertices = [
+ {"label": "person", "properties": {"name": "Tom Hanks", "age":
67}}, # age: INT -> int
+ {"label": "movie", "properties": {"title": "Forrest Gump", "year":
1994}}, # year: INT -> int
+ ]
+
+ edges = [
+ {
+ "label": "acted_in",
+ "properties": {"role": "Forrest Gump"}, # role: TEXT -> str
+ "outV": "person:Tom Hanks",
+ "inV": "movie:Forrest Gump",
+ }
+ ]
+
+ # Call the method - should succeed with correct data types
+ self.commit2graph.load_into_graph(vertices, edges, self.schema)
+
+ # Verify that _handle_graph_creation was called for each vertex and
edge
+ self.assertEqual(mock_handle_graph_creation.call_count, 3) # 2
vertices + 1 edge
+
+
@patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation")
+ def test_load_into_graph_maps_llm_vertex_ids_to_created_vertex_ids(self,
mock_handle_graph_creation):
+ """Test edges use server-created vertex ids when LLM ids differ."""
+ mock_handle_graph_creation.side_effect = [
+ MagicMock(id="1:Tom Hanks"),
+ MagicMock(id="2:Forrest Gump"),
+ MagicMock(id="edge_id"),
+ ]
+
+ vertices = [
+ {
+ "id": "person:Tom Hanks",
+ "label": "person",
+ "properties": {"name": "Tom Hanks", "age": 67},
+ },
Review Comment:
‼️ This test covers `CUSTOMIZE_STRING` only by calling `load_into_graph()`
directly, so it bypasses the real import flow that first calls
`init_schema_if_need()`. That production path still creates every vertex label
with `usePrimaryKeyId().primaryKeys(...)`, while `load_into_graph()` later
writes explicit ids for `id_strategy == "CUSTOMIZE_STRING"`. An inline schema
declaring custom string ids can therefore create a primary-key schema before
the custom-id write path runs. Please make `init_schema_if_need()` branch on
`id_strategy` and call `useCustomizeStringId()` for `CUSTOMIZE_STRING`, and add
an end-to-end `Commit2Graph.run()` or import-flow test for this schema.
##########
hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/commit_to_hugegraph.py:
##########
@@ -27,31 +27,109 @@
class Commit2Graph:
- def __init__(self):
+ def __init__(self, graph_config=None):
+ graph_config = graph_config or {}
+
+ def pick(key, default):
+ value = graph_config[key] if key in graph_config else default
+ return default if value is None else value
+
self.client = PyHugeClient(
- url=huge_settings.graph_url,
- graph=huge_settings.graph_name,
- user=huge_settings.graph_user,
- pwd=huge_settings.graph_pwd,
- graphspace=huge_settings.graph_space,
+ url=pick("url", huge_settings.graph_url),
+ graph=pick("graph", huge_settings.graph_name),
+ user=pick("user", huge_settings.graph_user),
+ pwd=pick("pwd", huge_settings.graph_pwd),
+ graphspace=pick("gs", huge_settings.graph_space),
)
self.schema = self.client.schema()
+ def _empty_import_result(self, vertices=None, edges=None, triples=None) ->
Dict[str, Any]:
+ return {
+ "vertices_attempted": len(vertices or []),
+ "vertices_created": 0,
+ "vertices_skipped": 0,
+ "edges_attempted": len(edges or []),
+ "edges_created": 0,
+ "edges_skipped": 0,
+ "triples_attempted": len(triples or []),
+ "triples_created": 0,
+ "triples_skipped": 0,
+ "errors": [],
+ }
+
+ def _import_error(self, kind, index, reason, label=None, key=None) ->
Dict[str, Any]:
+ error = {"kind": kind, "index": index, "reason": reason}
+ if label:
+ error["label"] = label
+ if key:
+ error["key"] = key
+ return error
+
+ def _validate_input_properties(
+ self,
+ kind,
+ index,
+ label,
+ properties,
+ allowed_properties,
+ property_label_map,
+ import_result,
+ ) -> bool:
+ allowed = set(allowed_properties)
+ skipped_key = "vertices_skipped" if kind == "vertex" else
"edges_skipped"
+ for key, value in properties.items():
+ property_label = property_label_map.get(key)
+ if key not in allowed or property_label is None:
+ log.error(
+ "(Input) %s property '%s' is not defined in schema label
'%s', skip it & need check it again",
+ kind,
+ key,
+ label,
+ )
+ import_result[skipped_key] += 1
+ import_result["errors"].append(self._import_error(kind, index,
"unknown_property", label, key))
+ return False
+ # TODO: transform to Enum first (better in earlier step)
+ data_type = property_label["data_type"]
+ cardinality = property_label["cardinality"]
+ if not self._check_property_data_type(data_type, cardinality,
value):
Review Comment:
⚠️ The operator-level property type validator does not match the public
request validator. The request model explicitly excludes `bool` for
`BYTE/INT/LONG`, but this path uses `isinstance(value, int)`, so Python accepts
`True`/`False` as integer values when named-schema or internal workflow paths
bypass the request-side inline-schema check. Please share one validator between
`GraphImportRequest` and `Commit2Graph`, and add coverage for `INT=True` plus
the intended `FLOAT/DOUBLE` behavior for JSON integer values.
--
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]
