This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hugegraph-ai.git
The following commit(s) were added to refs/heads/main by this push:
new eb8af86f Add POST /graph/extract REST API for programmatic graph
extraction. (#351)
eb8af86f is described below
commit eb8af86fa2c1401af58b1f45116b7139c394c7d9
Author: Nishita Matlani <[email protected]>
AuthorDate: Mon Jun 8 02:24:32 2026 -0400
Add POST /graph/extract REST API for programmatic graph extraction. (#351)
## Summary
Closes #348.
HugeGraph-LLM already supports graph extraction through the Gradio demo,
but there was no public REST endpoint for it. This PR adds `POST
/graph/extract` to the existing FastAPI app, routing requests through
`SchedulerSingleton` and `FlowName.GRAPH_EXTRACT` — the same path the
demo uses.
### Key changes
- Add `GraphExtractRequest` with validation for `texts`, `schema`,
`split_type`, and related options
- Add `graph_http_api` and register it on the existing auth router
- Make `split_type` configurable in `GraphExtractFlow` (default
`"document"`, so demo behavior is unchanged)
- Return structured JSON (`vertices` / `edges` as arrays), with optional
`warning` and `meta`
---------
Co-authored-by: Cursor <[email protected]>
Co-authored-by: imbajin <[email protected]>
---
.../src/hugegraph_llm/api/graph_extract_api.py | 69 ++++
.../api/models/graph_extract_requests.py | 127 +++++++
.../api/models/graph_extract_responses.py | 27 ++
.../src/hugegraph_llm/api/models/rag_requests.py | 2 +-
.../src/hugegraph_llm/demo/rag_demo/app.py | 2 +
.../src/hugegraph_llm/flows/graph_extract.py | 13 +
.../hugegraph_llm/nodes/hugegraph_node/schema.py | 2 +-
.../operators/hugegraph_op/schema_manager.py | 22 +-
hugegraph-llm/src/hugegraph_llm/state/ai_state.py | 3 +
.../src/tests/api/test_graph_extract_api.py | 396 +++++++++++++++++++++
.../src/tests/document/test_vector_index_utils.py | 5 +-
.../operators/hugegraph_op/test_schema_manager.py | 66 ++++
12 files changed, 726 insertions(+), 8 deletions(-)
diff --git a/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py
b/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py
new file mode 100644
index 00000000..6884412d
--- /dev/null
+++ b/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py
@@ -0,0 +1,69 @@
+# 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
+
+from fastapi import APIRouter, HTTPException, status
+
+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.utils.log import log
+
+
+class GraphExtractService:
+ @staticmethod
+ def extract_sync(req: GraphExtractRequest) -> GraphExtractResponse:
+ try:
+ scheduler = SchedulerSingleton.get_instance()
+ result_str = scheduler.schedule_flow(
+ FlowName.GRAPH_EXTRACT,
+ req.graph_schema,
+ req.texts,
+ req.example_prompt or prompt.extract_graph_prompt,
+ req.extract_type,
+ language=req.language,
+ split_type=req.split_type,
+ client_config=req.client_config,
+ )
+ raw = json.loads(result_str)
+ warnings = [raw.pop("warning")] if "warning" in raw else []
+ result = {"vertices": raw.get("vertices", []), "edges":
raw.get("edges", [])}
+ meta = {}
+ if req.include_meta:
+ meta = {
+ "vertex_count": len(result["vertices"]),
+ "edge_count": len(result["edges"]),
+ "text_count": len(req.texts),
+ }
+ return GraphExtractResponse(result=result, warnings=warnings,
meta=meta)
+ except HTTPException:
+ raise
+ except Exception as e:
+ log.error("Error in graph_extract_api: %s", e)
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail="An unexpected error occurred during graph extraction.",
+ ) from e
+
+
+def graph_extract_http_api(router: APIRouter):
+ @router.post("/graph/extract", status_code=status.HTTP_200_OK,
response_model=GraphExtractResponse)
+ def graph_extract_api(req: GraphExtractRequest):
+ return GraphExtractService.extract_sync(req)
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
new file mode 100644
index 00000000..d3e2654a
--- /dev/null
+++ b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_requests.py
@@ -0,0 +1,127 @@
+# 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
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from fastapi import Query
+from pydantic import BaseModel, ConfigDict, Field, field_validator,
model_validator
+
+
+class GraphExtractClientConfig(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+
+ graph: Optional[str] = None
+ user: Optional[str] = None
+ pwd: Optional[str] = None
+ gs: Optional[str] = None
+
+
+class GraphExtractRequest(BaseModel):
+ model_config = ConfigDict(populate_by_name=True)
+
+ texts: Union[str, List[str]] = Field(..., description="Text or list of
texts to extract a graph from.")
+ graph_schema: Union[str, Dict[str, Any]] = Field(
+ ...,
+ alias="schema",
+ description="Graph schema as a JSON string/object, or an existing
graph name.",
+ )
+ example_prompt: Optional[str] = Query(None, description="Optional graph
extraction prompt header.")
+ 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.")
+ 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.")
+
+ @field_validator("texts")
+ @classmethod
+ def normalize_texts(cls, v):
+ items = [v] if isinstance(v, str) else list(v)
+ items = [t for t in items if t and t.strip()]
+ if not items:
+ raise ValueError("texts must not be empty.")
+ return items
+
+ @field_validator("graph_schema")
+ @classmethod
+ def normalize_schema(cls, v):
+ def validate_schema_obj(schema_obj: Any) -> None:
+ if not isinstance(schema_obj, dict):
+ raise ValueError("schema JSON must be an object.")
+ if "vertexlabels" not in schema_obj or "edgelabels" not in
schema_obj:
+ raise ValueError("schema must contain 'vertexlabels' and
'edgelabels'.")
+ if not isinstance(schema_obj["vertexlabels"], list) or not
isinstance(schema_obj["edgelabels"], list):
+ raise ValueError("'vertexlabels' and 'edgelabels' must be
lists.")
+
+ for vlabel in schema_obj["vertexlabels"]:
+ if not isinstance(vlabel, dict):
+ raise ValueError("Each item in 'vertexlabels' must be an
object.")
+ if not isinstance(vlabel.get("name"), str) or not
vlabel["name"].strip():
+ raise ValueError("Each vertex label must have a non-empty
string 'name'.")
+ props = vlabel.get("properties")
+ if not isinstance(props, list) or len(props) == 0:
+ raise ValueError("Each vertex label must have a non-empty
'properties' list.")
+
+ for elabel in schema_obj["edgelabels"]:
+ if not isinstance(elabel, dict):
+ raise ValueError("Each item in 'edgelabels' must be an
object.")
+ for key in ("name", "source_label", "target_label"):
+ if not isinstance(elabel.get(key), str) or not
elabel[key].strip():
+ raise ValueError(f"Each edge label must have a
non-empty string '{key}'.")
+ if "properties" in elabel and not
isinstance(elabel["properties"], list):
+ raise ValueError("'properties' in edge labels must be a
list when provided.")
+
+ if "propertykeys" in schema_obj and not
isinstance(schema_obj["propertykeys"], list):
+ raise ValueError("'propertykeys' must be a list when
provided.")
+
+ if isinstance(v, dict):
+ validate_schema_obj(v)
+ return json.dumps(v, ensure_ascii=False)
+ v = v.strip()
+ if not v:
+ raise ValueError("schema must not be empty.")
+ if v.startswith("{"):
+ try:
+ schema_obj = json.loads(v)
+ except json.JSONDecodeError as e:
+ raise ValueError(f"Invalid JSON schema: {e}") from e
+ validate_schema_obj(schema_obj)
+ return v
+ return v
+
+ @model_validator(mode="after")
+ def validate_schema_and_client_config(self):
+ schema = self.graph_schema
+ is_named_schema = isinstance(schema, str) and not
schema.strip().startswith("{")
+ if not is_named_schema:
+ if self.client_config is not None:
+ raise ValueError(
+ "client_config is not allowed when 'schema' is inline
JSON; graph extraction "
+ "from an inline schema does not connect to HugeGraph."
+ )
+ return self
+ if self.client_config is None:
+ raise ValueError(
+ "client_config is required when 'schema' refers to an existing
graph name; "
+ "provide inline schema JSON instead to extract without a
HugeGraph connection."
+ )
+ if self.client_config.graph != schema:
+ raise ValueError(
+ "When 'schema' is a graph name, client_config.graph must match
it "
+ f"(got schema='{schema}',
client_config.graph='{self.client_config.graph}')."
+ )
+ return self
diff --git
a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_responses.py
b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_responses.py
new file mode 100644
index 00000000..25eab197
--- /dev/null
+++ b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_responses.py
@@ -0,0 +1,27 @@
+# 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 typing import Any, Dict, List, Literal
+
+from pydantic import BaseModel, Field
+
+
+class GraphExtractResponse(BaseModel):
+ status: Literal["succeeded"] = "succeeded"
+ result: Dict[str, Any]
+ warnings: List[str] = Field(default_factory=list)
+ meta: Dict[str, Any] = Field(default_factory=dict)
diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
index 77155978..df65e268 100644
--- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
+++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
@@ -32,7 +32,7 @@ class GraphConfigRequest(BaseModel):
graph: str = Query("hugegraph", description="hugegraph client name.")
user: str = Query("", description="hugegraph client user.")
pwd: str = Query("", description="hugegraph client pwd.")
- gs: str = None
+ gs: Optional[str] = None
class RAGRequest(BaseModel):
diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py
b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py
index c687df6d..fcb7cac2 100644
--- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py
+++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py
@@ -24,6 +24,7 @@ from fastapi import APIRouter, Depends, FastAPI
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from hugegraph_llm.api.admin_api import admin_http_api
+from hugegraph_llm.api.graph_extract_api import graph_extract_http_api
from hugegraph_llm.api.rag_api import rag_http_api
from hugegraph_llm.config import admin_settings, huge_settings, prompt
from hugegraph_llm.demo.rag_demo.admin_block import create_admin_block,
log_stream
@@ -178,6 +179,7 @@ def create_app():
gremlin_generate_selective,
)
admin_http_api(api_auth, log_stream)
+ graph_extract_http_api(api_auth)
app.include_router(api_auth)
# Mount Gradio inside FastAPI
diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py
b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py
index 13629e2b..4c96434a 100644
--- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py
+++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py
@@ -17,6 +17,7 @@ import json
from pycgraph import GPipeline
+from hugegraph_llm.config import huge_settings
from hugegraph_llm.flows.common import BaseFlow
from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode
from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode
@@ -55,6 +56,17 @@ class GraphExtractFlow(BaseFlow):
prepared_input.example_prompt = example_prompt
prepared_input.schema = schema
prepared_input.extract_type = extract_type
+ client_config = kwargs.get("client_config")
+ if client_config:
+ # URL stays server-controlled; only identity/graphspace are
request-scoped.
+ prepared_input.graph_client_config = {
+ "url": huge_settings.graph_url,
+ "user": client_config.user,
+ "pwd": client_config.pwd,
+ "graphspace": client_config.gs,
+ }
+ else:
+ prepared_input.graph_client_config = None
def build_flow(
self,
@@ -77,6 +89,7 @@ class GraphExtractFlow(BaseFlow):
extract_type,
split_type,
language,
+ **kwargs,
)
pipeline.createGParam(prepared_input, "wkflow_input")
diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py
b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py
index e9f9c608..c2c659b6 100644
--- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py
+++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py
@@ -39,7 +39,7 @@ class SchemaNode(BaseNode):
from_user_defined=None,
):
if from_hugegraph:
- return SchemaManager(from_hugegraph)
+ return SchemaManager(from_hugegraph,
connection=self.wk_input.graph_client_config)
if from_user_defined:
return CheckSchema(from_user_defined)
if from_extraction:
diff --git
a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/schema_manager.py
b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/schema_manager.py
index c51cb958..a862ea44 100644
--- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/schema_manager.py
+++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/schema_manager.py
@@ -23,14 +23,26 @@ from hugegraph_llm.config import huge_settings
class SchemaManager:
- def __init__(self, graph_name: str):
+ def __init__(self, graph_name: str, *, connection: Optional[Dict[str,
Any]] = None):
self.graph_name = graph_name
+ # Apply a request-scoped connection as a complete unit (omitted fields
stay as
+ # given) so it cannot fall back to global huge_settings per-field.
+ if connection is not None:
+ url = connection.get("url")
+ user = connection.get("user")
+ pwd = connection.get("pwd")
+ graphspace = connection.get("graphspace")
+ else:
+ url = huge_settings.graph_url
+ user = huge_settings.graph_user
+ pwd = huge_settings.graph_pwd
+ graphspace = huge_settings.graph_space
self.client = PyHugeClient(
- url=huge_settings.graph_url,
+ url=url,
graph=self.graph_name,
- user=huge_settings.graph_user,
- pwd=huge_settings.graph_pwd,
- graphspace=huge_settings.graph_space,
+ user=user,
+ pwd=pwd,
+ graphspace=graphspace,
)
self.schema = self.client.schema()
diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py
b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py
index 7dce7e85..739588c5 100644
--- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py
+++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py
@@ -26,6 +26,8 @@ class WkFlowInput(GParam):
split_type: Optional[str] = None # split type used by ChunkSplit Node
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.
+ graph_client_config: Optional[Dict[str, Any]] = None
data_json: Optional[Dict[str, Any]] = None
extract_type: Optional[str] = None
query_examples: Optional[Any] = None
@@ -86,6 +88,7 @@ class WkFlowInput(GParam):
self.split_type = None
self.example_prompt = None
self.schema = None
+ self.graph_client_config = None
self.data_json = None
self.extract_type = None
self.query_examples = None
diff --git a/hugegraph-llm/src/tests/api/test_graph_extract_api.py
b/hugegraph-llm/src/tests/api/test_graph_extract_api.py
new file mode 100644
index 00000000..29a3a827
--- /dev/null
+++ b/hugegraph-llm/src/tests/api/test_graph_extract_api.py
@@ -0,0 +1,396 @@
+# 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
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastapi import APIRouter, FastAPI, HTTPException, status
+from fastapi.testclient import TestClient
+from pydantic import ValidationError
+
+from hugegraph_llm.api.graph_extract_api import GraphExtractService,
graph_extract_http_api
+from hugegraph_llm.api.models.graph_extract_requests import
GraphExtractClientConfig, GraphExtractRequest
+from hugegraph_llm.api.models.graph_extract_responses import
GraphExtractResponse
+from hugegraph_llm.api.rag_api import rag_http_api
+from hugegraph_llm.config import huge_settings
+from hugegraph_llm.flows.graph_extract import GraphExtractFlow
+from hugegraph_llm.state.ai_state import WkFlowInput
+
+INLINE_SCHEMA = {"vertexlabels": [], "edgelabels": []}
+
+
+class CapturePipeline:
+ def __init__(self):
+ self.params = {}
+
+ def createGParam(self, value, name):
+ self.params[name] = value
+
+ def registerGElement(self, *args):
+ return None
+
+
+def _graph_client():
+ router = APIRouter()
+ graph_extract_http_api(router)
+ app = FastAPI()
+ app.include_router(router)
+ return TestClient(app)
+
+
+def _named_client_config(graph="custom_graph"):
+ return {"graph": graph, "user": "admin", "pwd": "secret", "gs": "space_a"}
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_returns_envelope(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps({"vertices": [{"id":
"1"}], "edges": []})
+ mock_singleton.get_instance.return_value = scheduler
+
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "张三在北京工作。", "schema": INLINE_SCHEMA, "include_meta":
True},
+ )
+
+ assert response.status_code == status.HTTP_200_OK
+ body = response.json()
+ assert body["status"] == "succeeded"
+ assert body["result"] == {"vertices": [{"id": "1"}], "edges": []}
+ assert body["warnings"] == []
+ assert body["meta"] == {"vertex_count": 1, "edge_count": 0, "text_count":
1}
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_omits_meta_by_default(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps({"vertices": [],
"edges": []})
+ mock_singleton.get_instance.return_value = scheduler
+
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": INLINE_SCHEMA})
+
+ assert response.status_code == status.HTTP_200_OK
+ assert response.json()["meta"] == {}
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_moves_warning_into_warnings(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps(
+ {"vertices": [], "edges": [], "warning": "The schema may not match the
Doc"}
+ )
+ mock_singleton.get_instance.return_value = scheduler
+
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": INLINE_SCHEMA})
+
+ body = response.json()
+ assert body["warnings"] == ["The schema may not match the Doc"]
+ assert "warning" not in body["result"]
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_accepts_text_and_list(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps({"vertices": [],
"edges": []})
+ mock_singleton.get_instance.return_value = scheduler
+
+ client = _graph_client()
+ client.post("/graph/extract", json={"texts": "single", "schema":
INLINE_SCHEMA})
+ assert scheduler.schedule_flow.call_args.args[2] == ["single"]
+
+ client.post("/graph/extract", json={"texts": ["a", "b"], "schema":
INLINE_SCHEMA})
+ assert scheduler.schedule_flow.call_args.args[2] == ["a", "b"]
+
+
+def test_graph_extract_rejects_empty_texts():
+ response = _graph_client().post("/graph/extract", json={"texts": " ",
"schema": INLINE_SCHEMA})
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_invalid_schema():
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": "{bad"})
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_incomplete_schema():
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": {"vertexlabels": []}})
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def
test_graph_extract_rejects_malformed_inline_schema_before_scheduler(mock_singleton):
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": {"vertexlabels": [{"name": "person"}],
"edgelabels": []}},
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+ mock_singleton.get_instance.assert_not_called()
+
+
+def test_graph_extract_rejects_invalid_split_type():
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": INLINE_SCHEMA, "split_type": "doc"},
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_triples_extract_type():
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": INLINE_SCHEMA, "extract_type":
"triples"},
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_named_schema_without_client_config():
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": "hugegraph"})
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_client_config_with_inline_schema():
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": INLINE_SCHEMA, "client_config":
_named_client_config()},
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_mismatched_schema_and_client_config_graph():
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": "custom_graph", "client_config":
_named_client_config("other_graph")},
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+def test_graph_extract_rejects_url_in_client_config():
+ response = _graph_client().post(
+ "/graph/extract",
+ json={
+ "texts": "x",
+ "schema": "custom_graph",
+ "client_config": {"graph": "custom_graph", "url": "10.0.0.1:8080"},
+ },
+ )
+ assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_named_schema_does_not_mutate_globals(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps({"vertices": [],
"edges": []})
+ mock_singleton.get_instance.return_value = scheduler
+
+ original = (
+ huge_settings.graph_url,
+ huge_settings.graph_name,
+ huge_settings.graph_user,
+ huge_settings.graph_pwd,
+ huge_settings.graph_space,
+ )
+ response = _graph_client().post(
+ "/graph/extract",
+ json={"texts": "x", "schema": "custom_graph", "client_config":
_named_client_config()},
+ )
+
+ assert response.status_code == status.HTTP_200_OK
+ assert (
+ huge_settings.graph_url,
+ huge_settings.graph_name,
+ huge_settings.graph_user,
+ huge_settings.graph_pwd,
+ huge_settings.graph_space,
+ ) == original
+ assert scheduler.schedule_flow.call_args.kwargs["client_config"].graph ==
"custom_graph"
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_graph_extract_scheduler_error_returns_500(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.side_effect = RuntimeError("Error in flow init")
+ mock_singleton.get_instance.return_value = scheduler
+
+ response = _graph_client().post("/graph/extract", json={"texts": "x",
"schema": INLINE_SCHEMA})
+ assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_service_extract_sync_builds_envelope(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.return_value = json.dumps({"vertices": [{"id":
"1"}], "edges": []})
+ mock_singleton.get_instance.return_value = scheduler
+
+ resp = GraphExtractService.extract_sync(GraphExtractRequest(texts="x",
schema=INLINE_SCHEMA, include_meta=True))
+
+ assert isinstance(resp, GraphExtractResponse)
+ assert resp.status == "succeeded"
+ assert resp.result == {"vertices": [{"id": "1"}], "edges": []}
+ assert resp.warnings == []
+ assert resp.meta == {"vertex_count": 1, "edge_count": 0, "text_count": 1}
+
+
+@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton")
+def test_service_extract_sync_maps_errors_to_500(mock_singleton):
+ scheduler = MagicMock()
+ scheduler.schedule_flow.side_effect = RuntimeError("boom")
+ mock_singleton.get_instance.return_value = scheduler
+
+ with pytest.raises(HTTPException) as exc_info:
+ GraphExtractService.extract_sync(GraphExtractRequest(texts="x",
schema=INLINE_SCHEMA))
+ assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
+
+
+def test_request_model_validation():
+ req = GraphExtractRequest(texts="hello", schema=INLINE_SCHEMA)
+ assert req.texts == ["hello"]
+ assert req.graph_schema == json.dumps(INLINE_SCHEMA, ensure_ascii=False)
+ assert req.client_config is None
+
+ with pytest.raises(ValidationError):
+ GraphExtractRequest(texts=[], schema="hugegraph")
+
+
+def test_request_model_named_schema_requires_matching_client_config():
+ with pytest.raises(ValidationError):
+ GraphExtractRequest(texts="hello", schema="hugegraph")
+
+ with pytest.raises(ValidationError):
+ GraphExtractRequest(
+ texts="hello",
+ schema="custom_graph",
+ client_config=GraphExtractClientConfig(graph="other_graph"),
+ )
+
+ req = GraphExtractRequest(
+ texts="hello",
+ schema="hugegraph",
+ client_config=GraphExtractClientConfig(graph="hugegraph",
user="admin", pwd="secret", gs="space_a"),
+ )
+ assert req.client_config.graph == "hugegraph"
+
+
+def test_request_model_rejects_client_config_with_inline_schema():
+ with pytest.raises(ValidationError):
+ GraphExtractRequest(
+ texts="hello",
+ schema=INLINE_SCHEMA,
+ client_config=GraphExtractClientConfig(graph="hugegraph"),
+ )
+
+
+def test_client_config_forbids_unknown_fields():
+ with pytest.raises(ValidationError):
+ GraphExtractClientConfig(graph="custom_graph", url="10.0.0.1:8080")
+
+
+def test_flow_prepare_sets_request_local_graph_config():
+ flow = GraphExtractFlow()
+ prepared_input = WkFlowInput()
+ client_config = GraphExtractClientConfig(graph="custom_graph",
user="admin", pwd="secret", gs="space_a")
+
+ flow.prepare(prepared_input, "custom_graph", ["text"], "prompt",
"property_graph", client_config=client_config)
+
+ assert prepared_input.graph_client_config == {
+ "url": huge_settings.graph_url,
+ "user": "admin",
+ "pwd": "secret",
+ "graphspace": "space_a",
+ }
+
+
+def test_flow_prepare_keeps_omitted_graphspace_none():
+ flow = GraphExtractFlow()
+ prepared_input = WkFlowInput()
+ client_config = GraphExtractClientConfig(graph="custom_graph",
user="admin", pwd="secret")
+
+ flow.prepare(prepared_input, "custom_graph", ["text"], "prompt",
"property_graph", client_config=client_config)
+
+ assert prepared_input.graph_client_config["graphspace"] is None
+
+
+def test_flow_prepare_does_not_leak_config_across_runs():
+ # A pooled pipeline is reused across requests, so prepare() must clear
config
+ # when a later request omits client_config.
+ flow = GraphExtractFlow()
+ prepared_input = WkFlowInput()
+ client_config = GraphExtractClientConfig(graph="custom_graph",
user="admin", pwd="secret", gs="space_a")
+
+ flow.prepare(prepared_input, "custom_graph", ["text"], "prompt",
"property_graph", client_config=client_config)
+ assert prepared_input.graph_client_config is not None
+
+ flow.prepare(prepared_input, "custom_graph", ["text"], "prompt",
"property_graph")
+ assert prepared_input.graph_client_config is None
+
+
+def test_flow_build_flow_preserves_split_type_and_client_config(monkeypatch):
+ monkeypatch.setattr(
+ "hugegraph_llm.flows.graph_extract.GPipeline",
+ CapturePipeline,
+ )
+ client_config = GraphExtractClientConfig(graph="custom_graph",
user="admin", pwd="secret", gs="space_a")
+
+ pipeline = GraphExtractFlow().build_flow(
+ "custom_graph",
+ ["text"],
+ "prompt",
+ "property_graph",
+ split_type="paragraph",
+ client_config=client_config,
+ )
+
+ prepared_input = pipeline.params["wkflow_input"]
+ assert prepared_input.split_type == "paragraph"
+ assert prepared_input.graph_client_config == {
+ "url": huge_settings.graph_url,
+ "user": "admin",
+ "pwd": "secret",
+ "graphspace": "space_a",
+ }
+
+
+def test_wkflow_input_reset_clears_graph_client_config():
+ prepared_input = WkFlowInput()
+ prepared_input.graph_client_config = {"url": "10.0.0.1:8080"}
+
+ prepared_input.reset(None)
+
+ assert prepared_input.graph_client_config is None
+
+
+def test_existing_routes_still_register():
+ router = APIRouter()
+ rag_http_api(
+ router,
+ rag_answer_func=Mock(),
+ graph_rag_recall_func=Mock(),
+ apply_graph_conf=Mock(),
+ apply_llm_conf=Mock(),
+ apply_embedding_conf=Mock(),
+ apply_reranker_conf=Mock(),
+ gremlin_generate_selective_func=Mock(),
+ )
+ graph_extract_http_api(router)
+ app = FastAPI()
+ app.include_router(router)
+
+ paths = {route.path for route in app.routes if hasattr(route, "path")}
+ assert "/rag" in paths
+ assert "/text2gremlin" in paths
+ assert "/config/graph" in paths
+ assert "/graph/extract" in paths
diff --git a/hugegraph-llm/src/tests/document/test_vector_index_utils.py
b/hugegraph-llm/src/tests/document/test_vector_index_utils.py
index 33213cdc..059121f0 100644
--- a/hugegraph-llm/src/tests/document/test_vector_index_utils.py
+++ b/hugegraph-llm/src/tests/document/test_vector_index_utils.py
@@ -192,9 +192,11 @@ def
test_read_documents_rejects_unsupported_file_type(tmp_path):
class DummyScheduler:
def __init__(self):
self.calls = []
+ self.kwargs = []
- def schedule_flow(self, *args):
+ def schedule_flow(self, *args, **kwargs):
self.calls.append(args)
+ self.kwargs.append(kwargs)
return "scheduled"
@@ -251,6 +253,7 @@ def
test_extract_graph_accepts_pdf_upload_and_forwards_text(monkeypatch, tmp_pat
assert "Extract Graph Entrypoint PDF" in texts[0]
assert forwarded_prompt == example_prompt
assert graph_mode == "property_graph"
+ assert scheduler.kwargs[0] == {"split_type": "document"}
def test_read_documents_rejects_encrypted_pdf(tmp_path):
diff --git
a/hugegraph-llm/src/tests/operators/hugegraph_op/test_schema_manager.py
b/hugegraph-llm/src/tests/operators/hugegraph_op/test_schema_manager.py
index 05d9659f..7ccf7310 100644
--- a/hugegraph-llm/src/tests/operators/hugegraph_op/test_schema_manager.py
+++ b/hugegraph-llm/src/tests/operators/hugegraph_op/test_schema_manager.py
@@ -91,6 +91,72 @@ class TestSchemaManager(unittest.TestCase):
self.assertEqual(self.schema_manager.client, self.mock_client)
self.assertEqual(self.schema_manager.schema, self.mock_schema)
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.PyHugeClient")
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.huge_settings")
+ def test_init_uses_request_local_connection_settings(self, mock_settings,
mock_client_class):
+ mock_settings.graph_url = "default:8080"
+ mock_settings.graph_user = "default_user"
+ mock_settings.graph_pwd = "default_pwd"
+ mock_settings.graph_space = "default_space"
+
+ SchemaManager(
+ "custom_graph",
+ connection={
+ "url": "10.0.0.1:8080",
+ "user": "admin",
+ "pwd": "secret",
+ "graphspace": "space_a",
+ },
+ )
+
+ mock_client_class.assert_called_once_with(
+ url="10.0.0.1:8080",
+ graph="custom_graph",
+ user="admin",
+ pwd="secret",
+ graphspace="space_a",
+ )
+
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.PyHugeClient")
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.huge_settings")
+ def test_init_request_config_does_not_inherit_global_graphspace(self,
mock_settings, mock_client_class):
+ mock_settings.graph_url = "default:8080"
+ mock_settings.graph_user = "default_user"
+ mock_settings.graph_pwd = "default_pwd"
+ mock_settings.graph_space = "global_space"
+
+ SchemaManager(
+ "custom_graph",
+ connection={
+ "url": "10.0.0.1:8080",
+ "user": "admin",
+ "pwd": "secret",
+ "graphspace": None,
+ },
+ )
+
+ _, kwargs = mock_client_class.call_args
+ assert kwargs["graphspace"] is None
+ assert kwargs["url"] == "10.0.0.1:8080"
+
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.PyHugeClient")
+ @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.huge_settings")
+ def test_init_falls_back_to_globals_without_connection(self,
mock_settings, mock_client_class):
+ mock_settings.graph_url = "default:8080"
+ mock_settings.graph_user = "default_user"
+ mock_settings.graph_pwd = "default_pwd"
+ mock_settings.graph_space = "global_space"
+
+ SchemaManager("custom_graph")
+
+ mock_client_class.assert_called_once_with(
+ url="default:8080",
+ graph="custom_graph",
+ user="default_user",
+ pwd="default_pwd",
+ graphspace="global_space",
+ )
+
def test_simple_schema_with_full_schema(self):
"""Test simple_schema method with a full schema."""
# Call the method
