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..a4f1046ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/graph_extract_api.py @@ -15,55 +15,355 @@ # 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" + + +def _validation_error_for_path(path: str, message: str) -> dict: + if path.endswith("/graph/import"): + return _error("GRAPH_IMPORT_VALIDATION_ERROR", message, "import") + if path.endswith("/graph/extract-and-import"): + return _error("GRAPH_EXTRACT_AND_IMPORT_VALIDATION_ERROR", message, "request") + return _error("GRAPH_EXTRACT_VALIDATION_ERROR", message, "request") + + +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: + message = _validation_message(exc.errors()) + return JSONResponse( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + content={"detail": _validation_error_for_path(request.url.path, message)}, + ) + + return custom_route_handler + + +def graph_extract_http_api( + router: APIRouter, + service=None, + job_store=None, + import_service=None, + run_jobs_inline: Optional[bool] = None, +): + extract_service = service or GraphExtractService() + graph_import_service = import_service or GraphImportService() + jobs = job_store or InMemoryGraphExtractJobStore() + original_route_class = router.route_class + router.route_class = GraphExtractAPIRoute + + @router.post("/graph/extract", status_code=status.HTTP_200_OK, response_model=GraphExtractResponse) + def graph_extract_api(req: GraphExtractRequest) -> GraphExtractResponse: + try: + return extract_service.extract_sync(req) + except FlowOutputValidationError as exc: + log.error("Graph extraction flow output is invalid: %s", exc) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_EXTRACT_INVALID_FLOW_OUTPUT", GRAPH_EXTRACT_FLOW_OUTPUT_ERROR, "extract"), + ) from exc + except ValueError as exc: + log.error("Graph extraction input is invalid: %s", exc) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error("GRAPH_EXTRACT_INVALID_INPUT", str(exc), "request"), + ) from exc + except Exception as exc: + log.error("Unexpected graph extraction error: %s", exc, exc_info=True) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_EXTRACT_FAILED", GRAPH_EXTRACT_RUNTIME_ERROR, "extract"), + ) from exc + + @router.post( + "/graph/extract/jobs", + status_code=status.HTTP_202_ACCEPTED, + response_model=GraphExtractJobCreateResponse, + ) + def create_graph_extract_job( + req: GraphExtractRequest, + ) -> GraphExtractJobCreateResponse: + """Create a process-local graph extraction job. + + Jobs are stored in memory and are not shared across API worker processes. + Job status/results are lost on service restart, and cancellation only applies before + a queued job starts running; it cannot interrupt an active LLM call. + """ + try: + job = jobs.create(req) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=_error("GRAPH_EXTRACT_JOB_LIMIT_EXCEEDED", str(exc), "job"), + ) from exc + if run_jobs_inline is True: + jobs.run_job(job.job_id, extract_service) + elif run_jobs_inline is None: + try: + jobs.submit_job(job.job_id, extract_service) + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=_error("GRAPH_EXTRACT_JOB_QUEUE_FULL", str(exc), "job", job.job_id), + ) from exc + return GraphExtractJobCreateResponse( + job_id=job.job_id, + status=job.status, + result_url=f"/graph/extract/jobs/{job.job_id}/result", + created_at=_job_ts(job.created_at), + updated_at=_job_ts(job.updated_at), + ) + + @router.get( + "/graph/extract/jobs/{job_id}", + status_code=status.HTTP_200_OK, + response_model=GraphExtractJobStatusResponse, + ) + def get_graph_extract_job(job_id: str) -> GraphExtractJobStatusResponse: + jobs.expire_jobs() + job = jobs.get(job_id) + if job is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=_error("GRAPH_EXTRACT_JOB_NOT_FOUND", f"Job {job_id} was not found", "job", job_id), + ) + return _job_status_response(job) + + @router.get( + "/graph/extract/jobs/{job_id}/result", + status_code=status.HTTP_200_OK, + response_model=GraphExtractResponse, + ) + def get_graph_extract_job_result(job_id: str) -> GraphExtractResponse: + jobs.expire_jobs() + job = jobs.get(job_id) + if job is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=_error("GRAPH_EXTRACT_JOB_NOT_FOUND", f"Job {job_id} was not found", "job", job_id), + ) + if job.status == GraphExtractJobStatus.EXPIRED: + raise HTTPException( + status_code=status.HTTP_410_GONE, + detail=_error("GRAPH_EXTRACT_JOB_EXPIRED", f"Job {job_id} result has expired", "job", job_id), + ) + if job.status in {GraphExtractJobStatus.PENDING, GraphExtractJobStatus.RUNNING}: + raise HTTPException( + status_code=status.HTTP_202_ACCEPTED, + detail=_error( + "GRAPH_EXTRACT_JOB_NOT_COMPLETE", + f"Job {job_id} is not complete", + "job", + job_id, + ), + ) + if job.status == GraphExtractJobStatus.CANCELLED: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=_error("GRAPH_EXTRACT_JOB_CANCELLED", f"Job {job_id} was cancelled", "job", job_id), + ) + if job.status == GraphExtractJobStatus.FAILED: + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=job.error.model_dump()) + if job.result is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error( + "GRAPH_EXTRACT_JOB_RESULT_MISSING", f"Job {job_id} finished without a result", "job", job_id + ), + ) + return job.result + + @router.delete( + "/graph/extract/jobs/{job_id}", + status_code=status.HTTP_200_OK, + response_model=GraphExtractJobStatusResponse, + ) + def cancel_graph_extract_job(job_id: str) -> GraphExtractJobStatusResponse: + job = jobs.cancel(job_id) + if job is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=_error("GRAPH_EXTRACT_JOB_NOT_FOUND", f"Job {job_id} was not found", "job", job_id), + ) + if job.status == GraphExtractJobStatus.RUNNING: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=_error( + "GRAPH_EXTRACT_JOB_NOT_CANCELLABLE", + f"Job {job_id} is already running and cannot be interrupted", + "job", + job_id, + ), + ) + return _job_status_response(job) + + @router.post("/graph/import", status_code=status.HTTP_200_OK, response_model=GraphImportResponse) + def graph_import_api(req: GraphImportRequest) -> GraphImportResponse: + if not req.write_to_graph: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error( + "GRAPH_IMPORT_CONFIRMATION_REQUIRED", + "write_to_graph=true is required before writing graph data to HugeGraph", + "import", + ), + ) 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, + return graph_import_service.import_graph(req) + except FlowOutputValidationError as exc: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_IMPORT_INVALID_FLOW_OUTPUT", GRAPH_IMPORT_FLOW_OUTPUT_ERROR, "import"), + ) from exc + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error("GRAPH_IMPORT_INVALID_INPUT", str(exc), "import"), + ) from exc + except Exception as exc: + log.error("Unexpected graph import error: %s", exc, exc_info=True) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_IMPORT_FAILED", GRAPH_IMPORT_RUNTIME_ERROR, "import"), + ) from exc + + @router.post( + "/graph/extract-and-import", + status_code=status.HTTP_200_OK, + response_model=GraphExtractAndImportResponse, + ) + def graph_extract_and_import_api(req: GraphExtractAndImportRequest) -> GraphExtractAndImportResponse: + if not req.write_to_graph: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error( + "GRAPH_IMPORT_CONFIRMATION_REQUIRED", + "write_to_graph=true is required before writing extraction results to HugeGraph", + "import", + ), ) - 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) + try: + extract_response = extract_service.extract_sync(req) + except FlowOutputValidationError as exc: + log.error("Extract-and-import extraction flow output is invalid: %s", exc) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail="An unexpected error occurred during graph extraction.", - ) from e + detail=_error("GRAPH_EXTRACT_INVALID_FLOW_OUTPUT", GRAPH_EXTRACT_FLOW_OUTPUT_ERROR, "extract"), + ) from exc + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error("GRAPH_EXTRACT_INVALID_INPUT", str(exc), "request"), + ) from exc + except Exception as exc: + log.error("Unexpected extract-and-import extraction error: %s", exc, exc_info=True) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_EXTRACT_FAILED", GRAPH_EXTRACT_RUNTIME_ERROR, "extract"), + ) from exc + try: + import_response = graph_import_service.import_graph( + GraphImportRequest( + schema=req.schema, + data=extract_response.result, + write_to_graph=True, + client_config=req.client_config, + options=req.import_options, + ) + ) + except FlowOutputValidationError as exc: + log.error("Extract-and-import import flow output is invalid: %s", exc) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_IMPORT_INVALID_FLOW_OUTPUT", GRAPH_IMPORT_FLOW_OUTPUT_ERROR, "import"), + ) from exc + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=_error("GRAPH_IMPORT_INVALID_INPUT", str(exc), "import"), + ) from exc + except Exception as exc: + log.error("Unexpected extract-and-import import error: %s", exc, exc_info=True) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=_error("GRAPH_IMPORT_FAILED", GRAPH_IMPORT_RUNTIME_ERROR, "import"), + ) from exc -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) + return GraphExtractAndImportResponse( + status=import_response.status, + extract_result=extract_response, + import_result=import_response, + ) + + router.route_class = original_route_class 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..0bab32da8 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 @@ -16,96 +16,189 @@ # under the License. import json +from copy import deepcopy from typing import Any, Dict, List, Literal, Optional, Union -from fastapi import Query -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_validator, model_validator + +from hugegraph_llm.config import llm_settings +from hugegraph_llm.operators.common_op.check_schema import CheckSchema +from hugegraph_llm.utils.schema_property import is_schema_property_value + +SchemaInput = Union[str, Dict[str, Any]] +ContentInput = Union[str, List[str]] +REQUIRED_VERTEX_KEYS = {"label", "properties"} +REQUIRED_EDGE_KEYS = {"label", "outV", "outVLabel", "inV", "inVLabel", "properties"} + + +def _validate_schema_value(schema: SchemaInput) -> SchemaInput: + if isinstance(schema, dict): + if not schema: + raise ValueError("schema must not be an empty object") + CheckSchema(deepcopy(schema)).run() + return schema + + schema_text = str(schema).strip() + if not schema_text: + raise ValueError("schema must not be empty") + 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 + if not isinstance(parsed_schema, dict) or not parsed_schema: + raise ValueError("schema JSON must be a non-empty object") + CheckSchema(deepcopy(parsed_schema)).run() + return schema_text + + +def _schema_object(schema: SchemaInput) -> Optional[Dict[str, Any]]: + if isinstance(schema, dict): + return schema + schema_text = str(schema).strip() + if not schema_text.startswith("{"): + return None + try: + parsed_schema = json.loads(schema_text) + except json.JSONDecodeError: + return None + return parsed_schema if isinstance(parsed_schema, dict) else None + + +class GraphExtractOptions(BaseModel): + include_meta: bool = Field(default=False, description="Whether to include response metadata.") + include_warnings: bool = Field(default=True, description="Whether to include extraction warnings.") + + +class GraphImportOptions(BaseModel): + update_vid_embeddings: bool = Field(default=False, description="Whether to rebuild vid embeddings after import.") class GraphExtractClientConfig(BaseModel): model_config = ConfigDict(extra="forbid") - graph: Optional[str] = None - user: Optional[str] = None - pwd: Optional[str] = None - gs: Optional[str] = None + graph: Optional[str] = Field(default=None, description="HugeGraph graph name.") + user: Optional[str] = Field(default=None, description="HugeGraph user.") + pwd: Optional[str] = Field(default=None, description="HugeGraph password.") + gs: Optional[str] = Field(default=None, description="HugeGraph graphspace.") + + @field_validator("graph", "user", "pwd", "gs", mode="before") + @classmethod + def blank_strings_to_none(cls, value): + if isinstance(value, str) and not value.strip(): + return None + return value 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( + content_type: Literal["text", "chunks"] = Field( + default="text", description="Whether content is raw text or chunks." + ) + content: Optional[ContentInput] = Field(default=None, description="Raw document text or pre-split chunks.") + texts: Optional[ContentInput] = Field(default=None, description="Deprecated alias for text or chunk content.") + schema_data: SchemaInput = Field( ..., alias="schema", - description="Graph schema as a JSON string/object, or an existing graph name.", + validation_alias=AliasChoices("schema", "graph_schema", "schema_data"), + serialization_alias="schema", + description="Graph schema JSON object/string, or 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()] + example_prompt: Optional[str] = Field(default=None, description="Extraction prompt header or examples.") + extract_type: Literal["property_graph"] = Field(default="property_graph") + language: Literal["zh", "en"] = Field(default="zh") + split_type: Literal["document", "paragraph", "sentence"] = Field(default="document") + max_parallel_chunks: Optional[int] = Field(default=None, description="Maximum chunk-level LLM calls per request.") + include_meta: bool = Field(default=False, description="Whether to include response metadata.") + include_warnings: bool = Field(default=True, description="Whether to include extraction warnings.") + client_config: Optional[GraphExtractClientConfig] = Field(default=None) + + @staticmethod + def _normalize_text_content(value: ContentInput) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError("content must be a non-empty string when content_type is text") + return value.strip() + + @staticmethod + def _normalize_legacy_texts(value: ContentInput) -> tuple[Literal["text", "chunks"], ContentInput]: + if isinstance(value, str): + return "text", GraphExtractRequest._normalize_text_content(value) + return "chunks", GraphExtractRequest._normalize_chunks(value) + + @staticmethod + def _normalize_chunks(value: ContentInput) -> List[str]: + if not isinstance(value, list): + raise ValueError("content must be a non-empty list of strings when content_type is chunks") + items = [] + for chunk in value: + if not isinstance(chunk, str) or not chunk.strip(): + raise ValueError("chunks content must contain only non-empty strings") + items.append(chunk.strip()) if not items: - raise ValueError("texts must not be empty.") + raise ValueError("chunks content must contain at least one non-empty string") return items - @field_validator("graph_schema") + @staticmethod + def _validate_parallel_chunks(value: Optional[int]) -> int: + requested = llm_settings.graph_extract_max_parallel_chunks if value is None else value + limit = llm_settings.graph_extract_max_parallel_chunks_limit + if requested < 1: + raise ValueError("max_parallel_chunks must be greater than or equal to 1") + if requested > limit: + raise ValueError(f"max_parallel_chunks must be less than or equal to {limit}") + return requested + + @property + def schema(self) -> SchemaInput: + return self.schema_data + + @property + def graph_schema(self) -> SchemaInput: + return self.schema_data + + @property + def options(self) -> GraphExtractOptions: + return GraphExtractOptions(include_meta=self.include_meta, include_warnings=self.include_warnings) + + @field_validator("schema_data") @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 + def validate_schema(cls, schema: SchemaInput) -> SchemaInput: + return _validate_schema_value(schema) @model_validator(mode="after") def validate_schema_and_client_config(self): - schema = self.graph_schema + if self.content is not None and self.texts is not None: + raise ValueError("content and deprecated texts alias cannot be provided together") + + if self.content is None: + if self.texts is None: + raise ValueError("content is required") + if "content_type" in self.model_fields_set: + raise ValueError("deprecated texts alias cannot be combined with content_type; use content instead") + legacy_content_type, legacy_content = self._normalize_legacy_texts(self.texts) + if legacy_content_type == "chunks" and self.split_type != "document": + raise ValueError("split_type must be 'document' when deprecated texts alias contains chunks") + self.content_type = legacy_content_type + self.content = legacy_content + self.texts = [legacy_content] if legacy_content_type == "text" else legacy_content + elif self.content_type == "text": + text = self._normalize_text_content(self.content) + self.content = text + self.texts = [text] + else: + if self.split_type != "document": + raise ValueError("split_type must be 'document' when content_type is chunks") + chunks = self._normalize_chunks(self.content) + self.content = chunks + self.texts = chunks + + self.max_parallel_chunks = self._validate_parallel_chunks(self.max_parallel_chunks) + + return self._validate_schema_client_config() + + def _validate_schema_client_config(self): + schema = self.schema_data is_named_schema = isinstance(schema, str) and not schema.strip().startswith("{") if not is_named_schema: if self.client_config is not None: @@ -125,3 +218,177 @@ def validate_schema_and_client_config(self): f"(got schema='{schema}', client_config.graph='{self.client_config.graph}')." ) return self + + +class GraphImportRequest(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + schema_data: SchemaInput = Field(..., alias="schema", description="Graph schema JSON object/string, or graph name.") + data: Dict[str, Any] = Field(..., description="Property graph data with vertices and edges.") + write_to_graph: bool = Field(default=False, description="Required confirmation for graph writes.") + client_config: Optional[GraphExtractClientConfig] = Field(default=None) + options: GraphImportOptions = Field(default_factory=GraphImportOptions) + + @property + def schema(self) -> SchemaInput: + return self.schema_data + + @field_validator("schema_data") + @classmethod + def validate_schema(cls, schema: SchemaInput) -> SchemaInput: + return _validate_schema_value(schema) + + @field_validator("data") + @classmethod + def validate_data(cls, data: Dict[str, Any]) -> Dict[str, Any]: + vertices = cls._optional_list(data, "vertices") + edges = cls._optional_list(data, "edges") + triples = cls._optional_list(data, "triples") + if triples: + raise ValueError("triples import is not supported; submit property graph vertices or edges") + if not vertices and not edges and not triples: + raise ValueError("data must contain at least one vertex or edge") + for index, vertex in enumerate(vertices): + if not isinstance(vertex, dict) or not REQUIRED_VERTEX_KEYS.issubset(vertex): + raise ValueError(f"vertices[{index}] must include label and properties") + if not isinstance(vertex["label"], str) or not vertex["label"].strip(): + raise ValueError(f"vertices[{index}].label must be a non-empty string") + if not isinstance(vertex["properties"], dict): + raise ValueError(f"vertices[{index}].properties must be an object") + for index, edge in enumerate(edges): + if not isinstance(edge, dict) or not REQUIRED_EDGE_KEYS.issubset(edge): + raise ValueError(f"edges[{index}] must include label, outV, outVLabel, inV, inVLabel, and properties") + for key in ("label", "outV", "outVLabel", "inV", "inVLabel"): + if not isinstance(edge[key], str) or not edge[key].strip(): + raise ValueError(f"edges[{index}].{key} must be a non-empty string") + if not isinstance(edge["properties"], dict): + raise ValueError(f"edges[{index}].properties must be an object") + return data + + @staticmethod + def _optional_list(data: Dict[str, Any], key: str) -> List[Any]: + if key not in data or data[key] is None: + return [] + if not isinstance(data[key], list): + raise ValueError(f"data.{key} must be a list") + return data[key] + + @model_validator(mode="after") + def validate_write_target(self): + self._validate_schema_properties() + self._validate_edge_endpoint_labels() + schema = self.schema_data + is_named_schema = isinstance(schema, str) and not schema.strip().startswith("{") + if ( + self.write_to_graph + and not is_named_schema + and (self.client_config is None or self.client_config.graph is None) + ): + raise ValueError("client_config.graph is required when writing inline schema data to HugeGraph") + if is_named_schema and self.client_config is not None and self.client_config.graph not in {None, schema}: + raise ValueError("schema graph name must match client_config.graph") + return self + + def _validate_schema_properties(self): + schema = _schema_object(self.schema_data) + if schema is None: + return + vertex_schema = { + vertex_label["name"]: vertex_label + for vertex_label in schema.get("vertexlabels", []) + if isinstance(vertex_label, dict) and "name" in vertex_label + } + edge_schema = { + edge_label["name"]: edge_label + for edge_label in schema.get("edgelabels", []) + if isinstance(edge_label, dict) and "name" in edge_label + } + property_schema = { + prop["name"]: prop for prop in schema.get("propertykeys", []) if isinstance(prop, dict) and "name" in prop + } + for index, vertex in enumerate(self.data.get("vertices", []) or []): + label = vertex.get("label") + schema_vertex = vertex_schema.get(label) + if schema_vertex is None: + raise ValueError(f"vertices[{index}].label is not defined in schema") + self._validate_item_properties( + f"vertices[{index}]", + str(label), + vertex.get("properties", {}), + schema_vertex.get("properties", []), + property_schema, + ) + for index, edge in enumerate(self.data.get("edges", []) or []): + label = edge.get("label") + schema_edge = edge_schema.get(label) + if schema_edge is None: + raise ValueError(f"edges[{index}].label is not defined in schema") + self._validate_item_properties( + f"edges[{index}]", + str(label), + edge.get("properties", {}), + schema_edge.get("properties", []), + property_schema, + ) + + @staticmethod + def _validate_item_properties( + item_path: str, + label: str, + properties: Dict[str, Any], + allowed_properties: List[str], + property_schema: Dict[str, Dict[str, Any]], + ): + allowed = set(allowed_properties) + for key, value in properties.items(): + if key not in allowed: + raise ValueError(f"{item_path}.properties.{key} is not defined for label '{label}'") + prop_schema = property_schema.get(key) + if prop_schema is not None and not is_schema_property_value(value, prop_schema): + raise ValueError(f"{item_path}.properties.{key} must match schema property type") + + def _validate_edge_endpoint_labels(self): + schema = _schema_object(self.schema_data) + if schema is None: + return + edge_schema = { + edge_label["name"]: edge_label + for edge_label in schema.get("edgelabels", []) + if isinstance(edge_label, dict) and "name" in edge_label + } + for index, edge in enumerate(self.data.get("edges", []) or []): + schema_edge = edge_schema.get(edge.get("label")) + if schema_edge is None: + continue + if edge.get("outVLabel") != schema_edge.get("source_label"): + raise ValueError( + f"edges[{index}].outVLabel must match schema source_label for edge label '{edge.get('label')}'" + ) + if edge.get("inVLabel") != schema_edge.get("target_label"): + raise ValueError( + f"edges[{index}].inVLabel must match schema target_label for edge label '{edge.get('label')}'" + ) + + +class GraphExtractAndImportRequest(GraphExtractRequest): + write_to_graph: bool = Field(default=False, description="Required confirmation for graph writes.") + import_options: GraphImportOptions = Field(default_factory=GraphImportOptions) + + def _validate_schema_client_config(self): + schema = self.schema_data + is_named_schema = isinstance(schema, str) and not schema.strip().startswith("{") + if not is_named_schema: + if self.write_to_graph and (self.client_config is None or self.client_config.graph is None): + raise ValueError("client_config.graph is required when writing inline schema data 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 index 25eab1973..7a9758d10 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_responses.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/graph_extract_responses.py @@ -15,13 +15,55 @@ # specific language governing permissions and limitations # under the License. -from typing import Any, Dict, List, Literal +from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, Field +class GraphExtractError(BaseModel): + code: str + message: str + phase: str + job_id: Optional[str] = None + + 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) + + +class GraphExtractJobCreateResponse(BaseModel): + job_id: str + status: str + result_url: str + created_at: str + updated_at: str + + +class GraphExtractJobStatusResponse(BaseModel): + job_id: str + status: str + created_at: str + updated_at: str + started_at: Optional[str] = None + finished_at: Optional[str] = None + expires_at: Optional[str] = None + error: Optional[GraphExtractError] = None + + +class GraphImportResponse(BaseModel): + status: str = "succeeded" + vertex_count: int = 0 + edge_count: int = 0 + triple_count: int = 0 + updated_embeddings: bool = False + warnings: List[str] = Field(default_factory=list) + meta: Dict[str, Any] = Field(default_factory=dict) + + +class GraphExtractAndImportResponse(BaseModel): + status: str = "succeeded" + extract_result: GraphExtractResponse + import_result: GraphImportResponse diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index fd2c82303..1c87facb6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -34,6 +34,8 @@ class LLMConfig(BaseConfig): keyword_extract_type: Literal["llm", "textrank", "hybrid"] = "llm" window_size: Optional[int] = 3 hybrid_llm_weights: Optional[float] = 0.5 + graph_extract_max_parallel_chunks: int = 2 + graph_extract_max_parallel_chunks_limit: int = 8 # TODO: divide RAG part if necessary # 1. OpenAI settings openai_chat_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") 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 fcb7cac2a..8aa9ff122 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -178,8 +178,8 @@ def create_app(): apply_reranker_config, gremlin_generate_selective, ) - admin_http_api(api_auth, log_stream) graph_extract_http_api(api_auth) + admin_http_api(api_auth, log_stream) 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 4c96434a6..50a4d6131 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -44,10 +44,21 @@ def prepare( extract_type, split_type=SPLIT_TYPE_DOCUMENT, language="zh", + content_type="text", + max_parallel_chunks=1, + client_config=None, **kwargs, ): # prepare input data prepared_input.texts = texts + if content_type not in {"text", "chunks"}: + raise ValueError("content_type must be text or chunks") + if content_type == "chunks" and split_type != SPLIT_TYPE_DOCUMENT: + raise ValueError("split_type must be document when content_type is chunks") + prepared_input.content_type = content_type + if not isinstance(max_parallel_chunks, int) or max_parallel_chunks < 1: + raise ValueError("max_parallel_chunks must be a positive integer") + prepared_input.max_parallel_chunks = max_parallel_chunks prepared_input.language = language if split_type not in VALID_SPLIT_TYPES: raise ValueError("split_type must be document, paragraph, or sentence") @@ -56,7 +67,6 @@ def prepare( 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 = { @@ -76,6 +86,9 @@ def build_flow( extract_type, split_type=SPLIT_TYPE_DOCUMENT, language="zh", + content_type="text", + max_parallel_chunks=1, + client_config=None, **kwargs, ): pipeline = GPipeline() @@ -87,9 +100,11 @@ def build_flow( texts, example_prompt, extract_type, - split_type, - language, - **kwargs, + split_type=split_type, + language=language, + content_type=content_type, + max_parallel_chunks=max_parallel_chunks, + client_config=client_config, ) pipeline.createGParam(prepared_input, "wkflow_input") @@ -110,19 +125,23 @@ def post_deal(self, pipeline=None, **kwargs): edges = res.get("edges", []) chunk_count = len(res.get("chunks", [])) log.info("Graph extraction chunk_count: %s", chunk_count) + output = { + "vertices": vertices, + "edges": edges, + "call_count": res.get("call_count"), + "chunk_count": chunk_count, + "max_parallel_chunks": res.get("max_parallel_chunks"), + } if not vertices and not edges: log.info("Please check the schema.(The schema may not match the Doc)") + output["warning"] = "The schema may not match the Doc" return json.dumps( - { - "vertices": vertices, - "edges": edges, - "warning": "The schema may not match the Doc", - }, + output, ensure_ascii=False, indent=2, ) return json.dumps( - {"vertices": vertices, "edges": edges}, + output, ensure_ascii=False, indent=2, ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py index ac0f2ab1a..7693e50f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -30,7 +30,7 @@ class ImportGraphDataFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, data, schema, **kwargs): + def prepare(self, prepared_input: WkFlowInput, data, schema, graph_config=None, **kwargs): try: data_json = json.loads(data.strip()) if isinstance(data, str) else data except json.JSONDecodeError as e: @@ -45,12 +45,13 @@ def prepare(self, prepared_input: WkFlowInput, data, schema, **kwargs): ) prepared_input.data_json = data_json prepared_input.schema = schema + prepared_input.graph_config = graph_config - def build_flow(self, data, schema, **kwargs): + def build_flow(self, data, schema, graph_config=None, **kwargs): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data - self.prepare(prepared_input, data, schema) + self.prepare(prepared_input, data, schema, graph_config=graph_config) pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py index 693f70012..f7eae71ef 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py @@ -23,14 +23,14 @@ # pylint: disable=arguments-differ,keyword-arg-before-vararg class UpdateVidEmbeddingsFlow(BaseFlow): - def prepare(self, prepared_input: WkFlowInput, **kwargs): - pass + def prepare(self, prepared_input: WkFlowInput, graph_config=None, **kwargs): + prepared_input.graph_config = graph_config - def build_flow(self, **kwargs): + def build_flow(self, graph_config=None, **kwargs): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data - self.prepare(prepared_input) + self.prepare(prepared_input, graph_config=graph_config) pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py index fafa1bdf4..e14108668 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -20,6 +20,100 @@ from hugegraph_llm.models.llms.ollama import OllamaClient from hugegraph_llm.models.llms.openai import OpenAIClient +OPENAI_DEFAULT_API_BASE = "https://api.openai.com/v1" +OPENAI_DEFAULT_MODEL = "gpt-4.1-mini" +OPENAI_DEFAULT_EXTRACT_TOKENS = 256 +LITELLM_DEFAULT_MODEL = "openai/gpt-4.1-mini" +LITELLM_DEFAULT_EXTRACT_TOKENS = 256 + + +def _extract_key_is_absent_or_shared(extract_api_key, chat_api_key) -> bool: + return not extract_api_key or extract_api_key == chat_api_key + + +def _use_openai_chat_fallback(llm_configs: LLMConfig) -> bool: + if llm_configs.chat_llm_type != "openai": + return False + explicit_extract_config = ( + not _extract_key_is_absent_or_shared(llm_configs.openai_extract_api_key, llm_configs.openai_chat_api_key) + or llm_configs.openai_extract_api_base not in {OPENAI_DEFAULT_API_BASE, llm_configs.openai_chat_api_base} + or ( + llm_configs.openai_extract_language_model + not in {OPENAI_DEFAULT_MODEL, llm_configs.openai_chat_language_model} + ) + or (llm_configs.openai_extract_tokens not in {OPENAI_DEFAULT_EXTRACT_TOKENS, llm_configs.openai_chat_tokens}) + ) + return not explicit_extract_config and bool( + llm_configs.openai_chat_api_key or llm_configs.openai_chat_api_base or llm_configs.openai_chat_language_model + ) + + +def _use_litellm_chat_fallback(llm_configs: LLMConfig) -> bool: + if llm_configs.chat_llm_type != "litellm": + return False + explicit_extract_config = ( + not _extract_key_is_absent_or_shared(llm_configs.litellm_extract_api_key, llm_configs.litellm_chat_api_key) + or llm_configs.litellm_extract_api_base not in {None, llm_configs.litellm_chat_api_base} + or ( + llm_configs.litellm_extract_language_model + not in {LITELLM_DEFAULT_MODEL, llm_configs.litellm_chat_language_model} + ) + or (llm_configs.litellm_extract_tokens not in {LITELLM_DEFAULT_EXTRACT_TOKENS, llm_configs.litellm_chat_tokens}) + ) + return not explicit_extract_config and bool( + llm_configs.litellm_chat_api_key or llm_configs.litellm_chat_api_base or llm_configs.litellm_chat_language_model + ) + + +def _ollama_extract_config(llm_configs: LLMConfig): + if ( + llm_configs.chat_llm_type == "ollama/local" + and not llm_configs.ollama_extract_language_model + and llm_configs.ollama_chat_language_model + ): + return { + "model": llm_configs.ollama_chat_language_model, + "host": llm_configs.ollama_chat_host, + "port": llm_configs.ollama_chat_port, + } + return { + "model": llm_configs.ollama_extract_language_model, + "host": llm_configs.ollama_extract_host, + "port": llm_configs.ollama_extract_port, + } + + +def _openai_extract_config(llm_configs: LLMConfig): + if _use_openai_chat_fallback(llm_configs): + return { + "api_key": llm_configs.openai_chat_api_key, + "api_base": llm_configs.openai_chat_api_base, + "model_name": llm_configs.openai_chat_language_model, + "max_tokens": llm_configs.openai_chat_tokens, + } + return { + "api_key": llm_configs.openai_extract_api_key, + "api_base": llm_configs.openai_extract_api_base, + "model_name": llm_configs.openai_extract_language_model, + "max_tokens": llm_configs.openai_extract_tokens, + } + + +def _litellm_extract_config(llm_configs: LLMConfig): + if _use_litellm_chat_fallback(llm_configs): + return { + "api_key": llm_configs.litellm_chat_api_key, + "api_base": llm_configs.litellm_chat_api_base, + "model_name": llm_configs.litellm_chat_language_model, + "max_tokens": llm_configs.litellm_chat_tokens, + } + return { + "api_key": llm_configs.litellm_extract_api_key, + "api_base": llm_configs.litellm_extract_api_base, + "model_name": llm_configs.litellm_extract_language_model, + "max_tokens": llm_configs.litellm_extract_tokens, + } + def get_chat_llm(llm_configs: LLMConfig): if llm_configs.chat_llm_type == "openai": @@ -47,25 +141,11 @@ def get_chat_llm(llm_configs: LLMConfig): def get_extract_llm(llm_configs: LLMConfig): if llm_configs.extract_llm_type == "openai": - return OpenAIClient( - api_key=llm_configs.openai_extract_api_key, - api_base=llm_configs.openai_extract_api_base, - model_name=llm_configs.openai_extract_language_model, - max_tokens=llm_configs.openai_extract_tokens, - ) + return OpenAIClient(**_openai_extract_config(llm_configs)) if llm_configs.extract_llm_type == "ollama/local": - return OllamaClient( - model=llm_configs.ollama_extract_language_model, - host=llm_configs.ollama_extract_host, - port=llm_configs.ollama_extract_port, - ) + return OllamaClient(**_ollama_extract_config(llm_configs)) if llm_configs.extract_llm_type == "litellm": - return LiteLLMClient( - api_key=llm_configs.litellm_extract_api_key, - api_base=llm_configs.litellm_extract_api_base, - model_name=llm_configs.litellm_extract_language_model, - max_tokens=llm_configs.litellm_extract_tokens, - ) + return LiteLLMClient(**_litellm_extract_config(llm_configs)) raise Exception("extract llm type is not supported !") @@ -124,25 +204,11 @@ def get_chat_llm(self): def get_extract_llm(self): if self.extract_llm_type == "openai": - return OpenAIClient( - api_key=llm_settings.openai_extract_api_key, - api_base=llm_settings.openai_extract_api_base, - model_name=llm_settings.openai_extract_language_model, - max_tokens=llm_settings.openai_extract_tokens, - ) + return OpenAIClient(**_openai_extract_config(llm_settings)) if self.extract_llm_type == "ollama/local": - return OllamaClient( - model=llm_settings.ollama_extract_language_model, - host=llm_settings.ollama_extract_host, - port=llm_settings.ollama_extract_port, - ) + return OllamaClient(**_ollama_extract_config(llm_settings)) if self.extract_llm_type == "litellm": - return LiteLLMClient( - api_key=llm_settings.litellm_extract_api_key, - api_base=llm_settings.litellm_extract_api_base, - model_name=llm_settings.litellm_extract_language_model, - max_tokens=llm_settings.litellm_extract_tokens, - ) + return LiteLLMClient(**_litellm_extract_config(llm_settings)) raise Exception("extract llm type is not supported !") def get_text2gql_llm(self): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py index 4da43e902..323725fc5 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -75,11 +75,11 @@ def run(self): try: res = self.operator_schedule(data_json) - except (ValueError, TypeError, KeyError, NotImplementedError) as exc: + except Exception as exc: err_msg = _format_node_err(self, exc) log.error(err_msg) - return CStatus(-1, err_msg) - # For unexpected exceptions, re-raise to let them propagate or be caught elsewhere + node_name = getattr(self, "name", None) or type(self).__name__ + return CStatus(-1, f"Node {node_name} failed: {type(exc).__name__}: {exc}") self.context.lock() try: diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py index 602c1cc00..44a7d9ee7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py @@ -33,8 +33,16 @@ def node_init(self): split_type = self.wk_input.split_type if isinstance(texts, str): texts = [texts] + if self.wk_input.content_type == "chunks": + self.chunk_split_op = None + return super().node_init() self.chunk_split_op = ChunkSplit(texts, split_type, language) return super().node_init() def operator_schedule(self, data_json): + if self.wk_input.content_type == "chunks": + context = data_json or {} + texts = self.wk_input.texts + context["chunks"] = texts if isinstance(texts, list) else [texts] + return context return self.chunk_split_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py index a4ebc7092..48206f3f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py @@ -27,7 +27,7 @@ def node_init(self): data_json = self.wk_input.data_json if self.wk_input.data_json else None if data_json: self.context.assign_from_json(data_json) - self.commit_to_graph_op = Commit2Graph() + self.commit_to_graph_op = Commit2Graph(self.wk_input.graph_config) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py index 6e9dd01ad..0660695cb 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py @@ -27,7 +27,7 @@ class FetchGraphDataNode(BaseNode): wk_input: Optional[WkFlowInput] = None def node_init(self): - client = get_hg_client() + client = get_hg_client(self.wk_input.graph_config if self.wk_input else None) self.fetch_graph_data_op = FetchGraphData(client) return super().node_init() 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 c2c659b62..758fe473d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -37,9 +37,12 @@ def _import_schema( from_hugegraph=None, from_extraction=None, from_user_defined=None, + graph_config=None, ): if from_hugegraph: - return SchemaManager(from_hugegraph, connection=self.wk_input.graph_client_config) + if self.wk_input.graph_client_config is not None: + return SchemaManager(from_hugegraph, connection=self.wk_input.graph_client_config) + return SchemaManager(from_hugegraph, graph_config=graph_config) if from_user_defined: return CheckSchema(from_user_defined) if from_extraction: @@ -59,7 +62,10 @@ def node_init(self): return CStatus(-1, f"Invalid JSON format in schema. {exc}") else: log.info("Get schema '%s' from graphdb.", self.schema) - self.schema_manager = self._import_schema(from_hugegraph=self.schema) + self.schema_manager = self._import_schema( + from_hugegraph=self.schema, + graph_config=self.wk_input.graph_config, + ) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py index 71c853720..e2fe7d6d9 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py @@ -32,7 +32,7 @@ def node_init(self): vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() - self.build_semantic_index_op = BuildSemanticIndex(embedding, vector_index) + self.build_semantic_index_op = BuildSemanticIndex(embedding, vector_index, self.wk_input.graph_config) return super().node_init() def operator_schedule(self, data_json): 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..b20a4eefb 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 @@ -16,7 +16,7 @@ from pycgraph import CStatus from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.models.llms.init_llm import get_extract_llm from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.info_extract import InfoExtract from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract @@ -32,7 +32,7 @@ class ExtractNode(BaseNode): extract_type: str = None def node_init(self): - llm = get_chat_llm(llm_settings) + llm = get_extract_llm(llm_settings) if self.wk_input.example_prompt is None: return CStatus(-1, "Error occurs when prepare for workflow input") example_prompt = self.wk_input.example_prompt @@ -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_parallel_chunks=self.wk_input.max_parallel_chunks or 1, + ) else: return CStatus(-1, f"Unsupported extract_type: {extract_type}") return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/chunk_split.py index 83a1d4bbe..d400c5cb3 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/chunk_split.py @@ -39,6 +39,10 @@ def _split_sentence_boundaries(text: str) -> list[str]: return [sentence.strip() for sentence in sentence_pattern.findall(text) if sentence.strip()] +def _split_paragraph_boundaries(text: str) -> list[str]: + return [paragraph.strip() for paragraph in re.split(r"\n\s*\n+", text) if paragraph.strip()] + + class ChunkSplit: def __init__( self, @@ -63,9 +67,16 @@ def _get_text_splitter(self, split_type: str): if split_type == SPLIT_TYPE_DOCUMENT: return lambda text: [text] if split_type == SPLIT_TYPE_PARAGRAPH: - return RecursiveCharacterTextSplitter( - chunk_size=500, chunk_overlap=30, separators=self.separators - ).split_text + text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=self.separators) + + def split_paragraphs(text: str) -> list[str]: + chunks = [] + paragraphs = _split_paragraph_boundaries(text) or [text] + for paragraph in paragraphs: + chunks.extend(text_splitter.split_text(paragraph)) + return chunks + + return split_paragraphs if split_type == SPLIT_TYPE_SENTENCE: return _split_sentence_boundaries raise ValueError("split_type must be document, paragraph, or sentence") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/commit_to_hugegraph.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/commit_to_hugegraph.py index 14bb38654..9d1d930ef 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/commit_to_hugegraph.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/commit_to_hugegraph.py @@ -24,34 +24,125 @@ from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType, default_value_map from hugegraph_llm.utils.log import log +from hugegraph_llm.utils.schema_property import is_property_value_for_type 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 _vertex_mapping_id(self, vertex, vertex_label=None): + explicit_id = vertex.get("id") + if explicit_id: + return explicit_id + if vertex_label is None: + return None + primary_keys = vertex_label.get("primary_keys", []) + properties = vertex.get("properties") or {} + if primary_keys and all(properties.get(pk) for pk in primary_keys): + return f"{vertex['label']}:{'!'.join(str(properties[pk]) for pk in primary_keys)}" + return None + + 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): + log.error( + "(Input) %s property type/format '%s' is not correct, skip it & need check it again", + kind, + key, + ) + import_result[skipped_key] += 1 + import_result["errors"].append(self._import_error(kind, index, "invalid_property_type", label, key)) + return False + return True + def run(self, data: dict) -> Dict[str, Any]: schema = data.get("schema") - vertices = data.get("vertices", []) - edges = data.get("edges", []) - if not vertices and not edges: - log.critical("(Loading) Both vertices and edges are empty. Please check the input data again.") - raise ValueError("Both vertices and edges input are empty.") + vertices = data.get("vertices", []) or [] + edges = data.get("edges", []) or [] + triples = data.get("triples", []) or [] + if not vertices and not edges and not triples: + log.critical("(Loading) vertices, edges, and triples are empty. Please check the input data again.") + raise ValueError("vertices, edges, and triples input are empty.") if not schema: - # TODO: ensure the function works correctly (update the logic later) - self.schema_free_mode(data.get("triples", [])) + if vertices or edges: + raise ValueError("Schema-free mode only supports triples input; vertices and edges require schema.") + if not triples: + raise ValueError("Schema-free mode requires non-empty triples input.") + import_result = self.schema_free_mode(triples) log.warning("Using schema_free mode, could try schema_define mode for better effect!") else: + if triples: + raise ValueError("Triples input is not supported when schema is provided; use vertices and edges.") + if not vertices and not edges: + log.critical( + "(Loading) property-graph vertices and edges are empty. Please check the input data again." + ) + raise ValueError("property-graph vertices and edges are required when schema is provided.") self.init_schema_if_need(schema) - self.load_into_graph(vertices, edges, schema) + import_result = self.load_into_graph(vertices, edges, schema) + data["import_result"] = import_result or self._empty_import_result(vertices, edges, triples) return data def _set_default_property(self, key, input_properties, property_label_map): @@ -78,12 +169,15 @@ def _handle_graph_creation(self, func, *args, **kwargs): def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many-statements # pylint: disable=R0912 (too-many-branches) + import_result = self._empty_import_result(vertices, edges) vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} edge_label_map = {e_label["name"]: e_label for e_label in schema["edgelabels"]} property_label_map = {p_label["name"]: p_label for p_label in schema["propertykeys"]} vid_mapping = {} # mapping from LLM-generated vertex ID to actual server vertex ID + batch_vertex_refs = set() + failed_vertex_refs = set() - for vertex in vertices: + for vertex_index, vertex in enumerate(vertices): input_label = vertex["label"] # 1. ensure the input_label in the graph schema if input_label not in vertex_label_map: @@ -91,6 +185,14 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- "(Input) VertexLabel %s not found in schema, skip & need check it!", input_label, ) + import_result["vertices_skipped"] += 1 + import_result["errors"].append( + self._import_error("vertex", vertex_index, "vertex_label_not_found", label=input_label) + ) + vertex_ref = self._vertex_mapping_id(vertex) + if vertex_ref: + batch_vertex_refs.add(vertex_ref) + failed_vertex_refs.add(vertex_ref) continue input_properties = vertex["properties"] @@ -110,6 +212,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex, ) has_problem = True + import_result["errors"].append( + self._import_error("vertex", vertex_index, "missing_primary_key", input_label, pk) + ) break # TODO: transform to Enum first (better in earlier step) data_type = property_label_map[pk]["data_type"] @@ -124,34 +229,37 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex, ) if has_problem: + import_result["vertices_skipped"] += 1 + vertex_ref = self._vertex_mapping_id(vertex, vertex_label) + if vertex_ref: + batch_vertex_refs.add(vertex_ref) + failed_vertex_refs.add(vertex_ref) continue + mapping_id = self._vertex_mapping_id(vertex, vertex_label) + if mapping_id: + batch_vertex_refs.add(mapping_id) # 3. Ensure all non-nullable props are set for key in non_null_keys: if key not in input_properties: self._set_default_property(key, input_properties, property_label_map) - # 4. Check all data type value is right - for key, value in input_properties.items(): - # TODO: transform to Enum first (better in earlier step) - data_type = property_label_map[key]["data_type"] - cardinality = property_label_map[key]["cardinality"] - if not self._check_property_data_type(data_type, cardinality, value): - log.error( - "Property type/format '%s' is not correct, skip it & need check it again", - key, - ) - has_problem = True - break - if has_problem: + # 4. Check all property keys and data types are schema-compliant. + if not self._validate_input_properties( + "vertex", + vertex_index, + input_label, + input_properties, + vertex_label["properties"], + property_label_map, + import_result, + ): + if mapping_id: + failed_vertex_refs.add(mapping_id) continue # TODO: we could try batch add vertices first, setback to single-mode if failed explicit_id = vertex.get("id") - mapping_id = explicit_id - if not mapping_id and primary_keys: - mapping_id = f"{input_label}:{'!'.join(str(input_properties[pk]) for pk in primary_keys)}" - if vertex_label.get("id_strategy") == "CUSTOMIZE_STRING" and explicit_id: result = self._handle_graph_creation( self.client.graph().addVertex, @@ -162,15 +270,20 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- else: result = self._handle_graph_creation(self.client.graph().addVertex, input_label, input_properties) if result is None: - raise ValueError(f"Failed to create vertex '{input_label}' with properties {input_properties}") + import_result["vertices_skipped"] += 1 + import_result["errors"].append( + self._import_error("vertex", vertex_index, "create_failed", label=input_label) + ) + if mapping_id: + failed_vertex_refs.add(mapping_id) + continue vid = result.id + import_result["vertices_created"] += 1 vertex["id"] = vid if mapping_id: vid_mapping[mapping_id] = vid - for edge in edges: - start = vid_mapping.get(edge.get("outV"), edge.get("outV")) - end = vid_mapping.get(edge.get("inV"), edge.get("inV")) + for edge_index, edge in enumerate(edges): label = edge["label"] properties = edge["properties"] @@ -179,10 +292,65 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- "(Input) EdgeLabel %s not found in schema, skip & need check it!", label, ) + import_result["edges_skipped"] += 1 + import_result["errors"].append( + self._import_error("edge", edge_index, "edge_label_not_found", label=label) + ) + continue + + edge_label = edge_label_map[label] + if edge.get("outVLabel") is not None and edge.get("outVLabel") != edge_label.get("source_label"): + import_result["edges_skipped"] += 1 + import_result["errors"].append( + self._import_error("edge", edge_index, "source_label_mismatch", label=label, key="outVLabel") + ) + continue + if edge.get("inVLabel") is not None and edge.get("inVLabel") != edge_label.get("target_label"): + import_result["edges_skipped"] += 1 + import_result["errors"].append( + self._import_error("edge", edge_index, "target_label_mismatch", label=label, key="inVLabel") + ) + continue + + start_ref = edge.get("outV") + end_ref = edge.get("inV") + endpoint_error = None + for endpoint_key, endpoint_ref in (("outV", start_ref), ("inV", end_ref)): + if endpoint_ref in failed_vertex_refs: + endpoint_error = ("endpoint_vertex_failed", endpoint_key) + break + if endpoint_ref in batch_vertex_refs and endpoint_ref not in vid_mapping: + endpoint_error = ("missing_endpoint", endpoint_key) + break + if endpoint_error is not None: + reason, endpoint_key = endpoint_error + import_result["edges_skipped"] += 1 + import_result["errors"].append( + self._import_error("edge", edge_index, reason, label=label, key=endpoint_key) + ) + continue + start = vid_mapping.get(start_ref, start_ref) + end = vid_mapping.get(end_ref, end_ref) + + if not self._validate_input_properties( + "edge", + edge_index, + label, + properties, + edge_label.get("properties", []), + property_label_map, + import_result, + ): continue # TODO: we could try batch add edges first, setback to single-mode if failed - self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) + result = self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) + if result is None: + import_result["edges_skipped"] += 1 + import_result["errors"].append(self._import_error("edge", edge_index, "create_failed", label=label)) + continue + import_result["edges_created"] += 1 + return import_result def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -197,9 +365,12 @@ def init_schema_if_need(self, schema: dict): properties = vertex["properties"] nullable_keys = vertex["nullable_keys"] primary_keys = vertex["primary_keys"] - self.schema.vertexLabel(vertex_label).properties(*properties).nullableKeys( - *nullable_keys - ).usePrimaryKeyId().primaryKeys(*primary_keys).ifNotExist().create() + vertex_builder = self.schema.vertexLabel(vertex_label).properties(*properties).nullableKeys(*nullable_keys) + if vertex.get("id_strategy") == "CUSTOMIZE_STRING": + vertex_builder = vertex_builder.useCustomizeStringId() + else: + vertex_builder = vertex_builder.usePrimaryKeyId().primaryKeys(*primary_keys) + vertex_builder.ifNotExist().create() for edge in edges: edge_label = edge["name"] @@ -211,6 +382,7 @@ def init_schema_if_need(self, schema: dict): ).properties(*properties).nullableKeys(*properties).ifNotExist().create() def schema_free_mode(self, data): + import_result = self._empty_import_result(triples=data) self.schema.propertyKey("name").asText().ifNotExist().create() self.schema.vertexLabel("vertex").useCustomizeStringId().properties("name").ifNotExist().create() self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( @@ -220,11 +392,23 @@ def schema_free_mode(self, data): self.schema.indexLabel("vertexByName").onV("vertex").by("name").secondary().ifNotExist().create() self.schema.indexLabel("edgeByName").onE("edge").by("name").secondary().ifNotExist().create() - for item in data: + for triple_index, item in enumerate(data): s, p, o = (element.strip() for element in item) - s_id = self.client.graph().addVertex("vertex", {"name": s}, id=s).id - t_id = self.client.graph().addVertex("vertex", {"name": o}, id=o).id - self.client.graph().addEdge("edge", s_id, t_id, {"name": p}) + s_vertex = self._handle_graph_creation(self.client.graph().addVertex, "vertex", {"name": s}, id=s) + t_vertex = self._handle_graph_creation(self.client.graph().addVertex, "vertex", {"name": o}, id=o) + if s_vertex is None or t_vertex is None: + import_result["triples_skipped"] += 1 + import_result["errors"].append(self._import_error("triple", triple_index, "create_vertices_failed")) + continue + edge = self._handle_graph_creation( + self.client.graph().addEdge, "edge", s_vertex.id, t_vertex.id, {"name": p} + ) + if edge is None: + import_result["triples_skipped"] += 1 + import_result["errors"].append(self._import_error("triple", triple_index, "create_edge_failed")) + continue + import_result["triples_created"] += 1 + return import_result def _create_property(self, prop: dict): name = prop["name"] @@ -284,12 +468,7 @@ def _set_property_cardinality(self, property_key, cardinality): log.error("Unknown cardinality %s for property_key %s", cardinality, property_key) def _check_property_data_type(self, data_type: str, cardinality: str, value) -> bool: - if cardinality in ( - PropertyCardinality.LIST.value, - PropertyCardinality.SET.value, - ): - return self._check_collection_data_type(data_type, value) - return self._check_single_data_type(data_type, value) + return is_property_value_for_type(data_type, cardinality, value, strict_data_type=True) def _check_collection_data_type(self, data_type: str, value) -> bool: if not isinstance(value, list): @@ -307,9 +486,9 @@ def _check_single_data_type(self, data_type: str, value) -> bool: PropertyDataType.INT.value, PropertyDataType.LONG.value, ): - return isinstance(value, int) + return isinstance(value, int) and not isinstance(value, bool) if data_type in (PropertyDataType.FLOAT.value, PropertyDataType.DOUBLE.value): - return isinstance(value, float) + return isinstance(value, (int, float)) and not isinstance(value, bool) if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below 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 a862ea446..6cc3700c2 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,20 +23,31 @@ class SchemaManager: - def __init__(self, graph_name: str, *, connection: Optional[Dict[str, Any]] = None): + def __init__( + self, + graph_name: str, + *, + connection: Optional[Dict[str, Any]] = None, + graph_config: 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. + + def pick(config, key, default): + value = config[key] if key in config else default + return default if value is None else value + if connection is not None: - url = connection.get("url") - user = connection.get("user") - pwd = connection.get("pwd") - graphspace = connection.get("graphspace") + connection = connection or {} + url = pick(connection, "url", huge_settings.graph_url) + user = pick(connection, "user", huge_settings.graph_user) + pwd = pick(connection, "pwd", huge_settings.graph_pwd) + graphspace = pick(connection, "graphspace", huge_settings.graph_space) else: - url = huge_settings.graph_url - user = huge_settings.graph_user - pwd = huge_settings.graph_pwd - graphspace = huge_settings.graph_space + graph_config = graph_config or {} + url = pick(graph_config, "url", huge_settings.graph_url) + user = pick(graph_config, "user", huge_settings.graph_user) + pwd = pick(graph_config, "pwd", huge_settings.graph_pwd) + graphspace = pick(graph_config, "gs", huge_settings.graph_space) self.client = PyHugeClient( url=url, graph=self.graph_name, diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py index 1c75ea4b6..852f3d9f6 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py @@ -28,10 +28,12 @@ class BuildSemanticIndex: - def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): - self.vid_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids") + def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase], graph_config=None): + graph_config = graph_config or {} + self.graph_name = graph_config.get("graph") or huge_settings.graph_name + self.vid_index = vector_index.from_name(embedding.get_embedding_dim(), self.graph_name, "graph_vids") self.embedding = embedding - self.sm = SchemaManager(huge_settings.graph_name) + self.sm = SchemaManager(self.graph_name, graph_config=graph_config) def _extract_names(self, vertices: list[str]) -> list[str]: return [v.split(":")[1] for v in vertices] @@ -72,7 +74,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: added_embeddings = asyncio.run(self._get_embeddings_parallel(vids_to_process)) log.info("Building vector index for %s vertices...", len(added_vids)) self.vid_index.add(added_embeddings, added_vids) - self.vid_index.save_index_by_name(huge_settings.graph_name, "graph_vids") + self.vid_index.save_index_by_name(self.graph_name, "graph_vids") else: log.debug("No update vertices to build vector index.") context.update( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/info_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/info_extract.py index a786e52d4..ca1639c97 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/info_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/info_extract.py @@ -95,6 +95,7 @@ def extract_triples_by_regex_with_schema(schema, text, graph): text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") pattern = r"\((.*?), (.*?), (.*?)\) - ([^ ]*)" matches = re.findall(pattern, text) + schema = _to_legacy_schema(schema) vertices_dict = {v["id"]: v for v in graph["vertices"]} for match in matches: @@ -148,6 +149,29 @@ def extract_triples_by_regex_with_schema(schema, text, graph): graph["vertices"] = list(vertices_dict.values()) +def _to_legacy_schema(schema): + if "vertices" in schema and "edges" in schema: + return schema + return { + "vertices": [ + { + "vertex_label": vertex["name"], + "properties": vertex.get("properties", []), + } + for vertex in schema.get("vertexlabels", []) + ], + "edges": [ + { + "edge_label": edge["name"], + "source_vertex_label": edge["source_label"], + "target_vertex_label": edge["target_label"], + "properties": edge.get("properties", []), + } + for edge in schema.get("edgelabels", []) + ], + } + + class InfoExtract: def __init__(self, llm: BaseLLM, example_prompt: Optional[str] = None) -> None: self.llm = llm 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..eae3b5795 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,6 +19,7 @@ import json import re +from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List from hugegraph_llm.config import prompt @@ -61,7 +62,7 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: "properties": vertex["properties"], } for edge in schema["edgelabels"]: - properties_map["edge"][edge["name"]] = {"properties": edge["properties"]} + properties_map["edge"][edge["name"]] = {"properties": edge.get("properties", [])} log.info("properties_map: %s", properties_map) for item in items: item_type = item["type"] @@ -78,9 +79,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_parallel_chunks: int = 1, + ) -> None: self.llm = llm self.example_prompt = example_prompt + self.max_parallel_chunks = max(1, max_parallel_chunks) self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: @@ -91,15 +98,33 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: if "edges" not in context: context["edges"] = [] items = [] - for chunk in chunks: - proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) + try: + max_parallel_chunks = max(1, int(context.get("max_parallel_chunks") or self.max_parallel_chunks)) + except (TypeError, ValueError): + max_parallel_chunks = max(1, self.max_parallel_chunks) + chunk_count = len(chunks) + if chunk_count == 0: + context["max_parallel_chunks"] = 1 + context["call_count"] = context.get("call_count", 0) + return context + worker_count = min(max_parallel_chunks, chunk_count) + context["max_parallel_chunks"] = worker_count + if worker_count <= 1: + proceeded_chunks = [self.extract_property_graph_by_llm(schema, chunk) for chunk in chunks] + else: + with ThreadPoolExecutor(max_workers=worker_count) as executor: + proceeded_chunks = list( + executor.map(lambda chunk: self.extract_property_graph_by_llm(schema, chunk), chunks) + ) + for index, (chunk, proceeded_chunk) in enumerate(zip(chunks, proceeded_chunks)): log.debug( - "[LLM] %s input: %s \n output:%s", + "[LLM] %s chunk processed: index=%s input_chars=%s output_chars=%s", self.__class__.__name__, - chunk, - proceeded_chunk, + index, + len(chunk), + len(proceeded_chunk) if isinstance(proceeded_chunk, str) else None, ) - items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) + items.extend(self._extract_and_filter_label(schema, proceeded_chunk, raise_on_invalid=True)) items = filter_item(schema, items) for item in items: if item["type"] == "vertex": @@ -165,6 +190,13 @@ def _resolve_endpoint(self, edge, endpoint_key, label_key, legacy_key, vertex_la label = legacy_endpoint.get("label") properties = legacy_endpoint.get("properties", {}) + if not isinstance(properties, dict): + log.warning( + "Invalid %s endpoint properties type '%s' has been ignored.", + legacy_key, + type(properties), + ) + return None, label if label not in vertex_label_map: return None, label canonical_id = self._primary_key_id(vertex_label_map[label], properties) @@ -191,10 +223,23 @@ def _normalize_edges(self, edges, edge_label_map, vertex_label_map, vertex_id_ma vertex_id_map, ) if not out_v or not in_v: - log.warning("Invalid edge endpoints '%s' have been ignored.", edge) + log.warning( + "Invalid edge endpoints have been ignored: label=%s, outVLabel=%s, inVLabel=%s.", + edge.get("label"), + out_v_label, + in_v_label, + ) continue if out_v_label != edge_label.get("source_label") or in_v_label != edge_label.get("target_label"): - log.warning("Invalid edge endpoint labels '%s' have been ignored.", edge) + log.warning( + "Invalid edge endpoint labels have been ignored: label=%s, outVLabel=%s, inVLabel=%s, " + "expectedOutVLabel=%s, expectedInVLabel=%s.", + edge.get("label"), + out_v_label, + in_v_label, + edge_label.get("source_label"), + edge_label.get("target_label"), + ) continue edge["outV"] = out_v @@ -204,7 +249,7 @@ 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]]: + def _extract_and_filter_label(self, schema, text, raise_on_invalid: bool = False) -> List[Dict[str, Any]]: # Strip markdown code blocks (e.g. ```json ... ```) text = re.sub(r"```\w*\n?", "", text) text = re.sub(r"```", "", text) @@ -214,6 +259,8 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: 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.") + if raise_on_invalid: + raise ValueError("Invalid property graph JSON: no JSON object or array found") return [] json_str = json_match.group(1).strip() @@ -228,6 +275,8 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: # 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'.") + if raise_on_invalid: + raise ValueError("Invalid property graph JSON: expecting 'vertices' and 'edges'") return items # Create sets for valid vertex and edge labels based on the schema @@ -248,6 +297,13 @@ def process_items(item_list, valid_labels, item_type): if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) continue + if not isinstance(item.get("properties"), dict): + log.warning( + "Invalid %s properties type '%s' has been ignored.", + item_type, + type(item.get("properties")), + ) + continue if item_type_value != item_type: log.warning("Invalid %s type '%s' has been ignored.", item_type, item_type_value) continue @@ -266,6 +322,8 @@ def process_items(item_list, valid_labels, item_type): 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: + except json.JSONDecodeError as exc: log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") + if raise_on_invalid: + raise ValueError("Invalid property graph JSON: failed to parse extracted JSON") from exc return items diff --git a/hugegraph-llm/src/hugegraph_llm/services/__init__.py b/hugegraph-llm/src/hugegraph_llm/services/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/services/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/hugegraph-llm/src/hugegraph_llm/services/graph_extract_jobs.py b/hugegraph-llm/src/hugegraph_llm/services/graph_extract_jobs.py new file mode 100644 index 000000000..842fff1a6 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/services/graph_extract_jobs.py @@ -0,0 +1,224 @@ +# 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 queue +import threading +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from enum import Enum +from typing import Any, Dict, List, Optional + +from hugegraph_llm.api.models.graph_extract_requests import GraphExtractRequest +from hugegraph_llm.api.models.graph_extract_responses import GraphExtractError, GraphExtractResponse +from hugegraph_llm.utils.log import log + +JOB_RUNTIME_ERROR = "Graph extraction job failed during execution" + + +class GraphExtractJobStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + EXPIRED = "expired" + + +@dataclass +class GraphExtractJob: + job_id: str + request: Any + status: GraphExtractJobStatus + created_at: datetime + updated_at: datetime + started_at: Optional[datetime] = None + finished_at: Optional[datetime] = None + expires_at: Optional[datetime] = None + result: Optional[GraphExtractResponse] = None + error: Optional[GraphExtractError] = None + + +class InMemoryGraphExtractJobStore: + def __init__(self, max_jobs: int = 100, result_ttl_seconds: int = 3600, max_running_jobs: int = 2): + self.max_jobs = max_jobs + self.result_ttl_seconds = result_ttl_seconds + self.max_running_jobs = max(1, max_running_jobs) + self._jobs: Dict[str, GraphExtractJob] = {} + self._lock = threading.RLock() + self._queue = queue.Queue(maxsize=max_jobs) + self._workers_started = False + + def create(self, request: Any) -> GraphExtractJob: + with self._lock: + self.cleanup() + if len(self._jobs) >= self.max_jobs: + raise ValueError("graph extraction job store is full") + now = self._now() + job = GraphExtractJob( + job_id=f"gex_{uuid.uuid4().hex}", + request=request, + status=GraphExtractJobStatus.PENDING, + created_at=now, + updated_at=now, + ) + self._jobs[job.job_id] = job + return job + + def get(self, job_id: str) -> Optional[GraphExtractJob]: + with self._lock: + return self._jobs.get(job_id) + + def list_jobs(self) -> List[GraphExtractJob]: + with self._lock: + return list(self._jobs.values()) + + def mark_running(self, job_id: str) -> Optional[GraphExtractJob]: + with self._lock: + job = self._jobs.get(job_id) + if job is None or job.status != GraphExtractJobStatus.PENDING: + return job + now = self._now() + job.status = GraphExtractJobStatus.RUNNING + job.started_at = now + job.updated_at = now + return job + + def mark_succeeded(self, job_id: str, result: GraphExtractResponse) -> Optional[GraphExtractJob]: + with self._lock: + job = self._jobs.get(job_id) + if job is None or job.status == GraphExtractJobStatus.CANCELLED: + return job + now = self._now() + job.status = GraphExtractJobStatus.SUCCEEDED + job.result = result + job.finished_at = now + job.expires_at = now + timedelta(seconds=self.result_ttl_seconds) + job.updated_at = now + return job + + def mark_failed(self, job_id: str, error: GraphExtractError) -> Optional[GraphExtractJob]: + with self._lock: + job = self._jobs.get(job_id) + if job is None or job.status == GraphExtractJobStatus.CANCELLED: + return job + now = self._now() + job.status = GraphExtractJobStatus.FAILED + job.error = error + job.finished_at = now + job.expires_at = now + timedelta(seconds=self.result_ttl_seconds) + job.updated_at = now + return job + + def cancel(self, job_id: str) -> Optional[GraphExtractJob]: + with self._lock: + job = self._jobs.get(job_id) + if job is None: + return None + if job.status == GraphExtractJobStatus.PENDING: + now = self._now() + job.status = GraphExtractJobStatus.CANCELLED + job.finished_at = now + job.expires_at = now + timedelta(seconds=self.result_ttl_seconds) + job.updated_at = now + return job + + def expire_jobs(self) -> None: + with self._lock: + self._expire_jobs_locked() + + def cleanup(self) -> None: + with self._lock: + self._expire_jobs_locked() + expired_job_ids = [ + job_id for job_id, job in self._jobs.items() if job.status == GraphExtractJobStatus.EXPIRED + ] + for job_id in expired_job_ids: + del self._jobs[job_id] + + def submit_job(self, job_id: str, service) -> Optional[GraphExtractJob]: + with self._lock: + job = self._jobs.get(job_id) + if job is None: + return None + if job.status != GraphExtractJobStatus.PENDING: + return job + self._start_workers_locked() + try: + self._queue.put_nowait((job_id, service)) + except queue.Full as exc: + self._jobs.pop(job_id, None) + raise ValueError("graph extraction job queue is full") from exc + return job + + def _expire_jobs_locked(self) -> None: + now = self._now() + for job in self._jobs.values(): + expired_by_result_ttl = job.expires_at is not None and job.expires_at <= now + expired_by_pending_ttl = ( + job.status == GraphExtractJobStatus.PENDING and self.result_ttl_seconds == 0 and job.created_at <= now + ) + if job.status != GraphExtractJobStatus.EXPIRED and (expired_by_result_ttl or expired_by_pending_ttl): + job.status = GraphExtractJobStatus.EXPIRED + job.result = None + job.updated_at = now + + def _start_workers_locked(self) -> None: + if self._workers_started: + return + for index in range(self.max_running_jobs): + thread = threading.Thread( + target=self._worker_loop, + name=f"graph-extract-job-worker-{index}", + daemon=True, + ) + thread.start() + self._workers_started = True + + def _worker_loop(self) -> None: + while True: + job_id, service = self._queue.get() + try: + self.run_job(job_id, service) + finally: + self._queue.task_done() + + def run_job(self, job_id: str, service) -> None: + job = self.mark_running(job_id) + if job is None or job.status != GraphExtractJobStatus.RUNNING: + return + try: + request = job.request + if isinstance(request, dict): + request = GraphExtractRequest(**request) + result = service.extract_sync(request) + self.mark_succeeded(job_id, result) + except Exception: # pylint: disable=broad-exception-caught + log.exception("Graph extraction job %s failed", job_id) + self.mark_failed( + job_id, + GraphExtractError( + code="GRAPH_EXTRACT_JOB_FAILED", + message=JOB_RUNTIME_ERROR, + phase="extract", + job_id=job_id, + ), + ) + + @staticmethod + def _now() -> datetime: + return datetime.now(timezone.utc) diff --git a/hugegraph-llm/src/hugegraph_llm/services/graph_extract_service.py b/hugegraph-llm/src/hugegraph_llm/services/graph_extract_service.py new file mode 100644 index 000000000..10e6ebda7 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/services/graph_extract_service.py @@ -0,0 +1,411 @@ +# 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 index, vertex in enumerate(vertices): + if not isinstance(vertex, dict) or "label" not in vertex or "properties" not in vertex: + raise FlowOutputValidationError("canonical property graph vertex must include label and properties") + if not _is_non_empty_string(vertex["label"]): + raise FlowOutputValidationError( + f"canonical property graph vertex[{index}].label must be a non-empty string" + ) + if not isinstance(vertex["properties"], dict): + raise FlowOutputValidationError(f"canonical property graph vertex[{index}].properties must be an object") + required_edge_keys = {"label", "outV", "outVLabel", "inV", "inVLabel", "properties"} + for index, edge in enumerate(edges): + if not isinstance(edge, dict) or not required_edge_keys.issubset(edge): + raise FlowOutputValidationError( + "canonical property graph edge must include label, outV, outVLabel, inV, inVLabel, and properties" + ) + for key in ("label", "outV", "outVLabel", "inV", "inVLabel"): + if not _is_non_empty_string(edge[key]): + raise FlowOutputValidationError( + f"canonical property graph edge[{index}].{key} must be a non-empty string" + ) + if not isinstance(edge["properties"], dict): + raise FlowOutputValidationError(f"canonical property graph edge[{index}].properties must be an object") + + +def _is_non_empty_string(value: Any) -> bool: + return isinstance(value, str) and bool(value.strip()) + + +def _build_import_status(import_result: Dict[str, Any]) -> str: + skipped = ( + import_result.get("vertices_skipped", 0) + + import_result.get("edges_skipped", 0) + + import_result.get("triples_skipped", 0) + ) + created = ( + import_result.get("vertices_created", 0) + + import_result.get("edges_created", 0) + + import_result.get("triples_created", 0) + ) + if skipped and created: + return "partial" + if skipped and not created: + return "failed" + return "succeeded" + + +def _validate_import_result(import_result: Any) -> Dict[str, Any]: + counter_keys = ( + "vertices_attempted", + "vertices_created", + "vertices_skipped", + "edges_attempted", + "edges_created", + "edges_skipped", + "triples_attempted", + "triples_created", + "triples_skipped", + ) + if not isinstance(import_result, dict): + raise FlowOutputValidationError("graph import flow output must include import_result") + for key in counter_keys: + value = import_result.get(key) + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise FlowOutputValidationError( + f"graph import flow output import_result.{key} must be a non-negative integer" + ) + if not isinstance(import_result.get("errors"), list): + raise FlowOutputValidationError("graph import flow output import_result.errors must be a list") + return import_result + + +def _sanitize_import_error(error: Any) -> Dict[str, Any]: + if not isinstance(error, dict): + return {"kind": "import", "reason": "import_error"} + sanitized = { + key: value for key, value in error.items() if key in SAFE_IMPORT_ERROR_KEYS and isinstance(value, (str, int)) + } + return sanitized or {"kind": "import", "reason": "import_error"} + + +def _format_import_warning(error: Dict[str, Any]) -> str: + kind = error.get("kind", "import") + reason = error.get("reason", "import_error") + parts = ["import error" if kind == "import" else f"{kind} import error"] + if "index" in error: + parts.append(f"index={error['index']}") + if "label" in error: + parts.append(f"label={error['label']}") + if "key" in error: + parts.append(f"key={error['key']}") + parts.append(f"reason={reason}") + return " ".join(parts) + + +class GraphExtractService: + def __init__(self, scheduler=None): + self._scheduler = scheduler + + @property + def scheduler(self): + return self._scheduler or SchedulerSingleton.get_instance() + + def extract_sync(self, request: GraphExtractRequest) -> GraphExtractResponse: + started = time.perf_counter() + schema = normalize_schema(request.schema) + extract_client_config = request.client_config if _schema_graph_name(schema) else None + client_config_meta = _redact_client_config(apply_client_config(extract_client_config, schema=schema)) + example_prompt = request.example_prompt or prompt.extract_graph_prompt + try: + raw_result = self.scheduler.schedule_flow( + FlowName.GRAPH_EXTRACT, + schema, + request.texts, + example_prompt, + request.extract_type, + language=request.language, + split_type=request.split_type, + client_config=extract_client_config, + content_type=request.content_type, + max_parallel_chunks=request.max_parallel_chunks, + ) + except Exception: + log.exception("Graph extraction failed during scheduler execution") + raise + + parsed_result = _parse_flow_json(raw_result, "Invalid graph extraction flow JSON") + warnings = _pop_warnings(parsed_result) + result = self._build_result(parsed_result, request.extract_type) + if not request.options.include_warnings: + warnings = [] + meta = ( + self._build_extract_meta(request, parsed_result, result, started, client_config_meta) + if request.options.include_meta + else {} + ) + return GraphExtractResponse(status="succeeded", result=result, warnings=warnings, meta=meta) + + def _build_result(self, parsed_result: Dict[str, Any], extract_type: str) -> Dict[str, Any]: + if extract_type == "triples": + triples = parsed_result.get("triples") + if not triples: + triples = self._legacy_edges_to_triples(parsed_result.get("edges", [])) + return {"triples": triples} + result = { + "vertices": parsed_result.get("vertices", []), + "edges": parsed_result.get("edges", []), + } + _validate_property_graph_result(result) + return result + + def _legacy_edges_to_triples(self, edges: Any) -> List[Dict[str, Any]]: + if not isinstance(edges, list): + return [] + triples = [] + for edge in edges: + if not isinstance(edge, dict): + continue + start = edge.get("start", edge.get("outV")) + end = edge.get("end", edge.get("inV")) + edge_type = edge.get("type", edge.get("label")) + if start is not None and end is not None and edge_type is not None: + triples.append({"start": start, "type": edge_type, "end": end}) + return triples + + def _build_extract_meta( + self, + request: GraphExtractRequest, + parsed_result: Dict[str, Any], + result: Dict[str, Any], + started: float, + client_config_meta: Dict[str, Any], + ) -> Dict[str, Any]: + chunk_count = parsed_result.get("chunk_count") + if chunk_count is None: + chunk_count = ( + len(request.texts) + if request.content_type == "chunks" and isinstance(request.texts, (list, tuple)) + else parsed_result.get("call_count") + ) + max_parallel_chunks = parsed_result.get("max_parallel_chunks") + if max_parallel_chunks is None: + max_parallel_chunks = ( + min(request.max_parallel_chunks, chunk_count) + if isinstance(chunk_count, int) and chunk_count >= 0 + else request.max_parallel_chunks + ) + + meta = { + "extract_type": request.extract_type, + "content_type": request.content_type, + "language": request.language, + "split_type": request.split_type, + "text_count": 1 if request.content_type == "text" else 0, + "chunk_count": chunk_count, + "max_parallel_chunks": max_parallel_chunks, + "vertex_count": _count_items(result, "vertices"), + "edge_count": _count_items(result, "edges"), + "triple_count": _count_items(result, "triples"), + "call_count": parsed_result.get("call_count"), + "duration_ms": int((time.perf_counter() - started) * 1000), + } + if client_config_meta: + meta["client_config"] = client_config_meta + return meta + + +class GraphImportService: + def __init__(self, scheduler=None): + self._scheduler = scheduler + + @property + def scheduler(self): + return self._scheduler or SchedulerSingleton.get_instance() + + def import_graph(self, request: GraphImportRequest) -> GraphImportResponse: + if not request.write_to_graph: + raise ValueError("write_to_graph must be True to confirm graph import") + + started = time.perf_counter() + schema = normalize_schema(request.schema) + graph_config = apply_client_config(request.client_config, schema=schema, align_graph_with_schema=True) + client_config_meta = _redact_client_config(graph_config) + try: + raw_result = self.scheduler.schedule_flow( + FlowName.IMPORT_GRAPH_DATA, + request.data, + schema, + graph_config=graph_config, + ) + except Exception: + log.exception("Graph import failed during scheduler execution") + raise + + parsed_result = _parse_flow_json(raw_result, "Invalid graph import flow JSON") + warnings = _pop_warnings(parsed_result) + import_result = _validate_import_result(parsed_result.get("import_result")) + import_errors = [_sanitize_import_error(item) for item in import_result.get("errors", [])] + import_result["errors"] = import_errors + warnings.extend(_format_import_warning(item) for item in import_errors) + status = _build_import_status(import_result) + vertex_count = int(import_result["vertices_created"]) + edge_count = int(import_result["edges_created"]) + triple_count = int(import_result["triples_created"]) + updated_embeddings = False + if request.options.update_vid_embeddings: + try: + self.scheduler.schedule_flow(FlowName.UPDATE_VID_EMBEDDINGS, graph_config=graph_config) + updated_embeddings = True + except Exception as exc: # pylint: disable=broad-exception-caught + log.warning("VID embedding update failed after graph import: %s", exc, exc_info=True) + warnings.append("update_vid_embeddings failed") + if status == "succeeded": + status = "partial" + + meta = {"duration_ms": int((time.perf_counter() - started) * 1000)} + meta["import_result"] = import_result + if client_config_meta: + meta["client_config"] = client_config_meta + return GraphImportResponse( + status=status, + vertex_count=vertex_count, + edge_count=edge_count, + triple_count=triple_count, + updated_embeddings=updated_embeddings, + warnings=warnings, + meta=meta, + ) + + def import_extracted_graph( + self, + request: GraphExtractAndImportRequest, + extract_response: GraphExtractResponse, + ) -> GraphImportResponse: + import_request = GraphImportRequest( + schema=request.schema, + data=extract_response.result, + write_to_graph=True, + client_config=request.client_config, + options=request.import_options, + ) + return self.import_graph(import_request) diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 739588c56..beb6fc867 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -22,6 +22,8 @@ class WkFlowInput(GParam): texts: Optional[Union[str, List[str]]] = None # texts input used by ChunkSplit Node + content_type: Optional[str] = None + max_parallel_chunks: Optional[int] = None language: Optional[str] = None # language configuration used by ChunkSplit Node split_type: Optional[str] = None # split type used by ChunkSplit Node example_prompt: Optional[str] = None # need by graph information extract @@ -29,6 +31,7 @@ class WkFlowInput(GParam): # 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 + graph_config: Optional[Dict[str, Any]] = None extract_type: Optional[str] = None query_examples: Optional[Any] = None few_shot_schema: Optional[Any] = None @@ -84,12 +87,15 @@ class WkFlowInput(GParam): def reset(self, _: CStatus) -> None: self.texts = None + self.content_type = None + self.max_parallel_chunks = None self.language = None self.split_type = None self.example_prompt = None self.schema = None self.graph_client_config = None self.data_json = None + self.graph_config = None self.extract_type = None self.query_examples = None self.few_shot_schema = None @@ -145,6 +151,7 @@ class WkFlowState(GParam): vertices: Optional[List[Any]] = None triples: Optional[List[Any]] = None call_count: Optional[int] = None + max_parallel_chunks: Optional[int] = None keywords: Optional[List[str]] = None vector_result: Optional[Any] = None @@ -203,6 +210,7 @@ def setup(self) -> CStatus: self.vertices = None self.triples = None self.call_count = None + self.max_parallel_chunks = None self.keywords = None self.vector_result = None diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index a151110f0..28ecad491 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -19,6 +19,7 @@ import os import shutil from datetime import datetime +from typing import Any, Mapping, Optional import requests from pyhugegraph.client import PyHugeClient @@ -38,13 +39,19 @@ def run_gremlin_query(query, fmt=True): return json.dumps(res, indent=4, ensure_ascii=False) if fmt else res -def get_hg_client(): +def get_hg_client(graph_config: Optional[Mapping[str, Any]] = None) -> PyHugeClient: + 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 + return 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), ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/schema_property.py b/hugegraph-llm/src/hugegraph_llm/utils/schema_property.py new file mode 100644 index 000000000..b992b06fd --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/utils/schema_property.py @@ -0,0 +1,74 @@ +# 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 re +from typing import Any, Dict + +COLLECTION_CARDINALITIES = {"LIST", "SET"} +INTEGER_DATA_TYPES = {"BYTE", "INT", "LONG"} +FLOAT_DATA_TYPES = {"FLOAT", "DOUBLE"} +TEXT_DATA_TYPES = {"TEXT", "UUID"} + + +def is_schema_property_value( + value: Any, + prop_schema: Dict[str, Any], + *, + strict_data_type: bool = False, +) -> bool: + data_type = prop_schema.get("data_type") + cardinality = prop_schema.get("cardinality") + if not data_type or not cardinality: + return True + return is_property_value_for_type( + data_type, + cardinality, + value, + strict_data_type=strict_data_type, + ) + + +def is_property_value_for_type( + data_type: str, + cardinality: str, + value: Any, + *, + strict_data_type: bool = False, +) -> bool: + cardinality = str(cardinality).upper() + if cardinality in COLLECTION_CARDINALITIES: + return isinstance(value, list) and all( + is_single_property_value(data_type, item, strict_data_type=strict_data_type) for item in value + ) + return is_single_property_value(data_type, value, strict_data_type=strict_data_type) + + +def is_single_property_value(data_type: str, value: Any, *, strict_data_type: bool = False) -> bool: + data_type = str(data_type).upper() + if data_type == "BOOLEAN": + return isinstance(value, bool) + if data_type in INTEGER_DATA_TYPES: + return isinstance(value, int) and not isinstance(value, bool) + if data_type in FLOAT_DATA_TYPES: + return isinstance(value, (int, float)) and not isinstance(value, bool) + if data_type in TEXT_DATA_TYPES: + return isinstance(value, str) + if data_type == "DATE": + return isinstance(value, str) and bool(re.match(r"^\d{4}-\d{2}-\d{2}$", value)) + if strict_data_type: + raise ValueError(f"Unknown/Unsupported data type: {data_type}") + return True 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..c797d92e3 100644 --- a/hugegraph-llm/src/tests/api/test_graph_extract_api.py +++ b/hugegraph-llm/src/tests/api/test_graph_extract_api.py @@ -16,22 +16,32 @@ # under the License. import json -from unittest.mock import MagicMock, Mock, patch +from concurrent.futures import ThreadPoolExecutor +from threading import Lock +from unittest.mock import Mock import pytest -from fastapi import APIRouter, FastAPI, HTTPException, status +from fastapi import APIRouter, FastAPI, 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.graph_extract_api import 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.config import huge_settings, llm_settings from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.services.graph_extract_service import ( + FlowOutputValidationError, + GraphExtractService, + normalize_schema, +) from hugegraph_llm.state.ai_state import WkFlowInput INLINE_SCHEMA = {"vertexlabels": [], "edgelabels": []} +VALID_SCHEMA = { + "vertexlabels": [{"name": "person", "properties": ["name"]}], + "edgelabels": [{"name": "knows", "source_label": "person", "target_label": "person"}], +} class CapturePipeline: @@ -45,9 +55,40 @@ def registerGElement(self, *args): return None -def _graph_client(): +class EchoGraphExtractService: + def __init__(self): + self.requests = [] + self._lock = Lock() + + def extract_sync(self, req): + with self._lock: + self.requests.append( + { + "content_type": req.content_type, + "texts": list(req.texts), + "max_parallel_chunks": req.max_parallel_chunks, + "client_config": req.client_config.model_dump() if req.client_config else None, + } + ) + chunk_count = len(req.texts) + return GraphExtractResponse( + status="succeeded", + result={"vertices": [], "edges": []}, + warnings=[], + meta={ + "content_type": req.content_type, + "chunk_count": chunk_count, + "max_parallel_chunks": min(req.max_parallel_chunks, chunk_count), + "call_count": chunk_count, + "texts": list(req.texts), + "client_config": req.client_config.model_dump() if req.client_config else None, + }, + ) + + +def _graph_client(service=None): router = APIRouter() - graph_extract_http_api(router) + graph_extract_http_api(router, service=service) app = FastAPI() app.include_router(router) return TestClient(app) @@ -57,146 +98,362 @@ 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 +def _graph_result(): + return { + "vertices": [{"label": "person", "properties": {"name": "marko"}}], + "edges": [ + { + "label": "knows", + "outV": "marko", + "outVLabel": "person", + "inV": "vadas", + "inVLabel": "person", + "properties": {}, + } + ], + } - response = _graph_client().post( + +def test_graph_extract_returns_envelope_from_service(): + service = Mock() + service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_graph_result(), + warnings=[], + meta={"vertex_count": 1, "edge_count": 1, "text_count": 1}, + ) + + response = _graph_client(service).post( "/graph/extract", - json={"texts": "张三在北京工作。", "schema": INLINE_SCHEMA, "include_meta": True}, + json={"texts": "marko knows vadas", "schema": VALID_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} + assert response.json() == { + "status": "succeeded", + "result": _graph_result(), + "warnings": [], + "meta": {"vertex_count": 1, "edge_count": 1, "text_count": 1}, + } + service.extract_sync.assert_called_once() -@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 +def test_graph_extract_accepts_content_text_wire_shape(): + service = Mock() + service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_graph_result(), + warnings=[], + meta={}, + ) - response = _graph_client().post("/graph/extract", json={"texts": "x", "schema": INLINE_SCHEMA}) + response = _graph_client(service).post( + "/graph/extract", + json={"content_type": "text", "content": "marko knows vadas", "schema": VALID_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"} + request = service.extract_sync.call_args.args[0] + assert request.content_type == "text" + assert request.content == "marko knows vadas" + assert request.texts == ["marko knows vadas"] + + +def test_graph_extract_accepts_content_chunks_wire_shape(): + service = Mock() + service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_graph_result(), + warnings=[], + meta={}, ) - 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"] + response = _graph_client(service).post( + "/graph/extract", + json={ + "content_type": "chunks", + "content": ["marko knows vadas", "vadas knows josh"], + "schema": VALID_SCHEMA, + "max_parallel_chunks": 2, + }, + ) -@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"] + assert response.status_code == status.HTTP_200_OK + request = service.extract_sync.call_args.args[0] + assert request.content_type == "chunks" + assert request.content == ["marko knows vadas", "vadas knows josh"] + assert request.texts == ["marko knows vadas", "vadas knows josh"] + assert request.max_parallel_chunks == 2 + + +def test_graph_extract_accepts_legacy_texts_list_as_chunks(): + service = Mock() + service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_graph_result(), + warnings=[], + meta={}, + ) - client.post("/graph/extract", json={"texts": ["a", "b"], "schema": INLINE_SCHEMA}) - assert scheduler.schedule_flow.call_args.args[2] == ["a", "b"] + response = _graph_client(service).post( + "/graph/extract", + json={"texts": ["marko knows vadas", "vadas knows josh"], "schema": VALID_SCHEMA}, + ) + assert response.status_code == status.HTTP_200_OK + request = service.extract_sync.call_args.args[0] + assert request.content_type == "chunks" + assert request.content == ["marko knows vadas", "vadas knows josh"] + assert request.texts == ["marko knows vadas", "vadas knows josh"] + + +def test_graph_extract_api_concurrent_requests_keep_request_state_isolated(): + service = EchoGraphExtractService() + + payloads = [ + { + "content_type": "text", + "content": "doc A paragraph one.\n\ndoc A paragraph two.", + "schema": VALID_SCHEMA, + "split_type": "paragraph", + "max_parallel_chunks": 2, + }, + { + "content_type": "chunks", + "content": ["doc B chunk one", "doc B chunk two"], + "schema": VALID_SCHEMA, + "max_parallel_chunks": 2, + }, + { + "texts": "legacy text alias", + "schema": "legacy_graph", + "client_config": _named_client_config("legacy_graph"), + }, + { + "content_type": "chunks", + "content": ["single direct chunk"], + "schema": VALID_SCHEMA, + "max_parallel_chunks": 4, + }, + ] + + def post_payload(payload): + return _graph_client(service).post("/graph/extract", json=payload).json() + + with ThreadPoolExecutor(max_workers=len(payloads)) as executor: + responses = list(executor.map(post_payload, payloads)) + + assert [response["status"] for response in responses] == ["succeeded"] * len(payloads) + assert responses[0]["meta"]["content_type"] == "text" + assert responses[0]["meta"]["texts"] == ["doc A paragraph one.\n\ndoc A paragraph two."] + assert responses[1]["meta"]["content_type"] == "chunks" + assert responses[1]["meta"]["texts"] == ["doc B chunk one", "doc B chunk two"] + assert responses[2]["meta"]["content_type"] == "text" + assert responses[2]["meta"]["client_config"]["graph"] == "legacy_graph" + assert responses[3]["meta"]["chunk_count"] == 1 + assert responses[3]["meta"]["max_parallel_chunks"] == 1 + assert len(service.requests) == len(payloads) + + +def test_graph_extract_rejects_invalid_public_contract_inputs(): + client = _graph_client(Mock()) + + cases = [ + {"texts": " ", "schema": INLINE_SCHEMA}, + {"texts": "x", "schema": "{bad"}, + {"texts": "x", "schema": {"vertexlabels": [{"name": "person"}], "edgelabels": []}}, + {"texts": "x", "schema": INLINE_SCHEMA, "split_type": "doc"}, + {"texts": "x", "schema": INLINE_SCHEMA, "extract_type": "triples"}, + {"content_type": "text", "content": ["chunk"], "schema": INLINE_SCHEMA}, + {"content_type": "chunks", "content": "not-a-list", "schema": INLINE_SCHEMA}, + {"content_type": "chunks", "content": [], "schema": INLINE_SCHEMA}, + {"content_type": "chunks", "content": ["x"], "schema": INLINE_SCHEMA, "split_type": "paragraph"}, + {"texts": "x", "content": "y", "schema": INLINE_SCHEMA}, + {"texts": "x", "schema": "hugegraph"}, + {"texts": "x", "schema": INLINE_SCHEMA, "client_config": _named_client_config()}, + {"texts": "x", "schema": "custom_graph", "client_config": _named_client_config("other_graph")}, + {"texts": "x", "schema": "custom_graph", "client_config": {"graph": "custom_graph", "url": "10.0.0.1:8080"}}, + ] + + for payload in cases: + response = client.post("/graph/extract", json=payload) + assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + + +def test_graph_extract_validation_error_reports_invalid_field_detail(): + response = _graph_client(Mock()).post( + "/graph/extract", + json={"content_type": "chunks", "content": ["chunk"], "schema": INLINE_SCHEMA, "split_type": "paragraph"}, + ) -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 + detail = response.json()["detail"] + assert detail["code"] == "GRAPH_EXTRACT_VALIDATION_ERROR" + assert detail["phase"] == "request" + error_detail = detail["message"] + assert "split_type" in error_detail + assert "document" in error_detail + assert "content_type is chunks" in error_detail + assert "input_value" not in error_detail + assert "input" not in error_detail + + +def test_graph_extract_validation_error_does_not_echo_sensitive_input(): + response = _graph_client(Mock()).post( + "/graph/extract", + json={ + "content_type": "text", + "content": "hello", + "schema": "custom_graph", + "client_config": {"graph": "custom_graph", "pwd": "top-secret", "url": "10.0.0.1:8080"}, + }, + ) - -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 + detail_text = json.dumps(response.json(), ensure_ascii=False) + assert "top-secret" not in detail_text + assert "10.0.0.1:8080" not in detail_text -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 - +def test_graph_extract_api_returns_structured_error_for_invalid_flow_output(): + service = Mock() + service.extract_sync.side_effect = FlowOutputValidationError( + "Invalid property graph JSON: failed to parse extracted JSON" + ) -@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton") -def test_graph_extract_rejects_malformed_inline_schema_before_scheduler(mock_singleton): - response = _graph_client().post( + response = _graph_client(service).post( "/graph/extract", - json={"texts": "x", "schema": {"vertexlabels": [{"name": "person"}], "edgelabels": []}}, + json={"content_type": "text", "content": "bad llm output", "schema": VALID_SCHEMA}, ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY - mock_singleton.get_instance.assert_not_called() + + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["detail"] == { + "code": "GRAPH_EXTRACT_INVALID_FLOW_OUTPUT", + "message": "Graph extraction flow output is invalid", + "phase": "extract", + } -def test_graph_extract_rejects_invalid_split_type(): - response = _graph_client().post( +def test_graph_extract_api_rejects_malformed_workflow_property_graph_output(): + scheduler = Mock() + scheduler.schedule_flow.return_value = json.dumps( + {"vertices": [{"label": "person", "properties": None}], "edges": []} + ) + + response = _graph_client(GraphExtractService(scheduler)).post( "/graph/extract", - json={"texts": "x", "schema": INLINE_SCHEMA, "split_type": "doc"}, + json={"content_type": "text", "content": "bad workflow output", "schema": VALID_SCHEMA}, ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["detail"] == { + "code": "GRAPH_EXTRACT_INVALID_FLOW_OUTPUT", + "message": "Graph extraction flow output is invalid", + "phase": "extract", + } -def test_graph_extract_rejects_triples_extract_type(): - response = _graph_client().post( + +def test_graph_extract_api_maps_client_value_error_to_bad_request(): + service = Mock() + service.extract_sync.side_effect = ValueError("schema graph name must match client_config.graph") + + response = _graph_client(service).post( "/graph/extract", - json={"texts": "x", "schema": INLINE_SCHEMA, "extract_type": "triples"}, + json={"content_type": "text", "content": "bad input", "schema": VALID_SCHEMA}, ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert response.json()["detail"] == { + "code": "GRAPH_EXTRACT_INVALID_INPUT", + "message": "schema graph name must match client_config.graph", + "phase": "request", + } -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_api_returns_structured_error_for_runtime_failure(): + service = Mock() + service.extract_sync.side_effect = RuntimeError("llm provider timeout") -def test_graph_extract_rejects_client_config_with_inline_schema(): - response = _graph_client().post( + response = _graph_client(service).post( "/graph/extract", - json={"texts": "x", "schema": INLINE_SCHEMA, "client_config": _named_client_config()}, + json={"content_type": "text", "content": "provider failure", "schema": VALID_SCHEMA}, ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["detail"] == { + "code": "GRAPH_EXTRACT_FAILED", + "message": "Graph extraction failed during execution", + "phase": "extract", + } + assert "llm provider timeout" not in json.dumps(response.json(), ensure_ascii=False) -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")}, + +def test_graph_extract_request_validates_content_shape_and_parallel_limit(monkeypatch): + monkeypatch.setattr(llm_settings, "graph_extract_max_parallel_chunks", 2) + monkeypatch.setattr(llm_settings, "graph_extract_max_parallel_chunks_limit", 3) + + text_request = GraphExtractRequest(content_type="text", content="hello", schema=INLINE_SCHEMA) + assert text_request.texts == ["hello"] + assert text_request.max_parallel_chunks == 2 + + chunk_request = GraphExtractRequest( + content_type="chunks", + content=["chunk one", "chunk two"], + schema=INLINE_SCHEMA, + max_parallel_chunks=3, ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert chunk_request.texts == ["chunk one", "chunk two"] + assert chunk_request.max_parallel_chunks == 3 + with pytest.raises(ValidationError): + GraphExtractRequest(content_type="chunks", content=["chunk"], schema=INLINE_SCHEMA, max_parallel_chunks=4) -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"}, - }, + +def test_graph_extract_service_parses_flow_json_and_records_metadata(): + scheduler = Mock() + scheduler.schedule_flow.return_value = json.dumps( + { + **_graph_result(), + "call_count": 2, + "chunk_count": 2, + "warning": "schema mismatch", + } ) - assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + response = GraphExtractService(scheduler).extract_sync( + GraphExtractRequest( + content_type="text", + content="marko knows vadas", + schema=VALID_SCHEMA, + language="en", + include_meta=True, + max_parallel_chunks=2, + ) + ) -@patch("hugegraph_llm.api.graph_extract_api.SchedulerSingleton") -def test_graph_extract_named_schema_does_not_mutate_globals(mock_singleton): - scheduler = MagicMock() + assert response.status == "succeeded" + assert response.result == _graph_result() + assert response.warnings == ["schema mismatch"] + assert response.meta["extract_type"] == "property_graph" + assert response.meta["language"] == "en" + assert response.meta["text_count"] == 1 + assert response.meta["content_type"] == "text" + assert response.meta["chunk_count"] == 2 + assert response.meta["max_parallel_chunks"] == 2 + assert response.meta["vertex_count"] == 1 + assert response.meta["edge_count"] == 1 + assert response.meta["call_count"] == 2 + scheduler.schedule_flow.assert_called_once() + assert scheduler.schedule_flow.call_args.kwargs["language"] == "en" + assert scheduler.schedule_flow.call_args.kwargs["split_type"] == "document" + assert scheduler.schedule_flow.call_args.kwargs["content_type"] == "text" + assert scheduler.schedule_flow.call_args.kwargs["max_parallel_chunks"] == 2 + + +def test_graph_extract_service_passes_request_local_client_config_and_redacts_password(monkeypatch): + scheduler = Mock() scheduler.schedule_flow.return_value = json.dumps({"vertices": [], "edges": []}) - mock_singleton.get_instance.return_value = scheduler - + monkeypatch.setattr(huge_settings, "graph_url", "127.0.0.1:8080") original = ( huge_settings.graph_url, huge_settings.graph_name, @@ -204,12 +461,25 @@ def test_graph_extract_named_schema_does_not_mutate_globals(mock_singleton): 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()}, + client_config = GraphExtractClientConfig(graph="custom_graph", user="admin", pwd="secret", gs="space_a") + + response = GraphExtractService(scheduler).extract_sync( + GraphExtractRequest( + texts="x", + schema="custom_graph", + client_config=client_config, + include_meta=True, + ) ) - assert response.status_code == status.HTTP_200_OK + assert scheduler.schedule_flow.call_args.kwargs["client_config"] == client_config + assert "graph_config" not in scheduler.schedule_flow.call_args.kwargs + assert response.meta["client_config"] == { + "graph": "custom_graph", + "user": "admin", + "pwd": "***", + "gs": "space_a", + } assert ( huge_settings.graph_url, huge_settings.graph_name, @@ -217,86 +487,60 @@ def test_graph_extract_named_schema_does_not_mutate_globals(mock_singleton): 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 +def test_graph_extract_service_rejects_invalid_flow_json(): + scheduler = Mock() + scheduler.schedule_flow.return_value = "{broken" - response = _graph_client().post("/graph/extract", json={"texts": "x", "schema": INLINE_SCHEMA}) - assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + with pytest.raises(ValueError, match="Invalid graph extraction flow JSON"): + GraphExtractService(scheduler).extract_sync(GraphExtractRequest(texts="x", schema=INLINE_SCHEMA)) -@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 +def test_normalize_schema_rejects_malformed_json_schema(): + with pytest.raises(ValueError, match="schema must be valid JSON"): + normalize_schema("{bad") - 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} +def test_property_graph_response_rejects_legacy_edge_shape(): + scheduler = Mock() + scheduler.schedule_flow.return_value = json.dumps( + { + "vertices": [{"label": "person", "properties": {"name": "marko"}}], + "edges": [{"start": "marko", "type": "knows", "end": "vadas"}], + } + ) + + with pytest.raises(ValueError, match="canonical property graph edge"): + GraphExtractService(scheduler).extract_sync(GraphExtractRequest(texts="x", schema=INLINE_SCHEMA)) -@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 +def test_property_graph_response_rejects_non_object_properties(): + scheduler = Mock() + scheduler.schedule_flow.return_value = json.dumps( + {"vertices": [{"label": "person", "properties": None}], "edges": []} + ) - 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 + with pytest.raises(ValueError, match=r"vertex\[0\].properties"): + GraphExtractService(scheduler).extract_sync(GraphExtractRequest(texts="x", schema=INLINE_SCHEMA)) -def test_request_model_validation(): +def test_request_model_validation_and_aliases(): req = GraphExtractRequest(texts="hello", schema=INLINE_SCHEMA) assert req.texts == ["hello"] - assert req.graph_schema == json.dumps(INLINE_SCHEMA, ensure_ascii=False) + assert req.graph_schema == INLINE_SCHEMA + assert req.schema == INLINE_SCHEMA 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"), - ) + legacy_schema_alias = GraphExtractRequest(texts="hello", graph_schema=INLINE_SCHEMA) + assert legacy_schema_alias.schema == INLINE_SCHEMA + legacy_chunks = GraphExtractRequest(texts=["hello", "world"], schema=INLINE_SCHEMA) + assert legacy_chunks.content_type == "chunks" + assert legacy_chunks.texts == ["hello", "world"] -def test_client_config_forbids_unknown_fields(): with pytest.raises(ValidationError): - GraphExtractClientConfig(graph="custom_graph", url="10.0.0.1:8080") + GraphExtractRequest(texts=[], schema="hugegraph") def test_flow_prepare_sets_request_local_graph_config(): @@ -314,35 +558,50 @@ def test_flow_prepare_sets_request_local_graph_config(): } -def test_flow_prepare_keeps_omitted_graphspace_none(): +def test_flow_prepare_rejects_invalid_max_parallel_chunks(): 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 + with pytest.raises(ValueError, match="max_parallel_chunks"): + flow.prepare(WkFlowInput(), INLINE_SCHEMA, ["text"], "prompt", "property_graph", max_parallel_chunks=0) -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. +def test_flow_prepare_preserves_content_type_and_parallel_chunks(): 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", + ["chunk one", "chunk two"], + "prompt", + "property_graph", + content_type="chunks", + max_parallel_chunks=3, + ) - flow.prepare(prepared_input, "custom_graph", ["text"], "prompt", "property_graph") - assert prepared_input.graph_client_config is None + assert prepared_input.texts == ["chunk one", "chunk two"] + assert prepared_input.content_type == "chunks" + assert prepared_input.max_parallel_chunks == 3 -def test_flow_build_flow_preserves_split_type_and_client_config(monkeypatch): - monkeypatch.setattr( - "hugegraph_llm.flows.graph_extract.GPipeline", - CapturePipeline, +def test_graph_extract_meta_handles_invalid_chunk_texts_defensively(): + request = GraphExtractRequest(content_type="chunks", content=["chunk"], schema=INLINE_SCHEMA, include_meta=True) + request.texts = None + + meta = GraphExtractService()._build_extract_meta( + request, + {"call_count": 1}, + {"vertices": [], "edges": []}, + 0, + {}, ) + + assert meta["chunk_count"] == 1 + assert meta["max_parallel_chunks"] == 1 + + +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( @@ -351,46 +610,22 @@ def test_flow_build_flow_preserves_split_type_and_client_config(monkeypatch): "prompt", "property_graph", split_type="paragraph", + max_parallel_chunks=3, 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", - } + assert prepared_input.max_parallel_chunks == 3 + assert prepared_input.graph_client_config["graphspace"] == "space_a" -def test_wkflow_input_reset_clears_graph_client_config(): +def test_wkflow_input_reset_clears_graph_configs(): prepared_input = WkFlowInput() prepared_input.graph_client_config = {"url": "10.0.0.1:8080"} + prepared_input.graph_config = {"graph": "custom_graph"} 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 = set(app.openapi()["paths"]) - assert "/rag" in paths - assert "/text2gremlin" in paths - assert "/config/graph" in paths - assert "/graph/extract" in paths + assert prepared_input.graph_config is None diff --git a/hugegraph-llm/src/tests/api/test_graph_extract_jobs.py b/hugegraph-llm/src/tests/api/test_graph_extract_jobs.py new file mode 100644 index 000000000..443c3a041 --- /dev/null +++ b/hugegraph-llm/src/tests/api/test_graph_extract_jobs.py @@ -0,0 +1,267 @@ +# 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 time +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): + router = APIRouter() + graph_extract_http_api( + router, + service=service, + job_store=job_store, + run_jobs_inline=run_jobs_inline, + ) + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def _success_response(): + return GraphExtractResponse( + status="succeeded", + result={"vertices": [{"label": "person"}], "edges": []}, + warnings=[], + meta={ + "extract_type": "property_graph", + "language": "zh", + "split_type": "document", + "text_count": 1, + "vertex_count": 1, + "edge_count": 0, + "call_count": 1, + "duration_ms": 1, + }, + ) + + +def test_job_creation_returns_pending_status_and_result_url_without_running_inline(): + client = _client(Mock(), InMemoryGraphExtractJobStore(), run_jobs_inline=False) + + response = client.post("/graph/extract/jobs", json=_payload()) + + assert response.status_code == status.HTTP_202_ACCEPTED + body = response.json() + assert body["job_id"] + assert body["status"] == GraphExtractJobStatus.PENDING + assert body["result_url"] == f"/graph/extract/jobs/{body['job_id']}/result" + + +def test_successful_job_reaches_succeeded_and_exposes_result(): + service = Mock() + service.extract_sync.return_value = _success_response() + store = InMemoryGraphExtractJobStore() + client = _client(service, store, run_jobs_inline=True) + + created = client.post("/graph/extract/jobs", json=_payload()).json() + status_response = client.get(f"/graph/extract/jobs/{created['job_id']}") + result_response = client.get(f"/graph/extract/jobs/{created['job_id']}/result") + + assert status_response.status_code == status.HTTP_200_OK + assert status_response.json()["status"] == GraphExtractJobStatus.SUCCEEDED + assert result_response.status_code == status.HTTP_200_OK + assert result_response.json()["result"]["vertices"] == [{"label": "person"}] + + +def test_default_job_route_runs_through_background_worker_and_exposes_result(): + service = Mock() + service.extract_sync.return_value = _success_response() + store = InMemoryGraphExtractJobStore() + client = _client(service, store, run_jobs_inline=None) + + created = client.post("/graph/extract/jobs", json=_payload()).json() + + status_body = {} + for _ in range(50): + status_response = client.get(f"/graph/extract/jobs/{created['job_id']}") + status_body = status_response.json() + if status_body["status"] == GraphExtractJobStatus.SUCCEEDED: + break + time.sleep(0.02) + + assert status_body["status"] == GraphExtractJobStatus.SUCCEEDED + result_response = client.get(f"/graph/extract/jobs/{created['job_id']}/result") + assert result_response.status_code == status.HTTP_200_OK + assert result_response.json()["status"] == "succeeded" + service.extract_sync.assert_called_once() + + +def test_failed_job_stores_error_details(): + service = Mock() + service.extract_sync.side_effect = RuntimeError("llm failed with sensitive raw document") + client = _client(service, InMemoryGraphExtractJobStore(), run_jobs_inline=True) + + created = client.post("/graph/extract/jobs", json=_payload()).json() + result_response = client.get(f"/graph/extract/jobs/{created['job_id']}/result") + + assert result_response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert result_response.json()["detail"]["code"] == "GRAPH_EXTRACT_JOB_FAILED" + assert result_response.json()["detail"]["message"] == "Graph extraction job failed during execution" + assert result_response.json()["detail"]["phase"] == "extract" + assert "sensitive raw document" not in str(result_response.json()) + + +def test_pending_job_result_returns_not_complete_semantics_and_can_be_cancelled(): + store = InMemoryGraphExtractJobStore() + client = _client(Mock(), store, run_jobs_inline=False) + + created = client.post("/graph/extract/jobs", json=_payload()).json() + pending_result = client.get(f"/graph/extract/jobs/{created['job_id']}/result") + delete_response = client.delete(f"/graph/extract/jobs/{created['job_id']}") + status_response = client.get(f"/graph/extract/jobs/{created['job_id']}") + + assert pending_result.status_code == status.HTTP_202_ACCEPTED + assert pending_result.json()["detail"]["code"] == "GRAPH_EXTRACT_JOB_NOT_COMPLETE" + assert delete_response.status_code == status.HTTP_200_OK + assert status_response.json()["status"] == GraphExtractJobStatus.CANCELLED + + +def test_cancelled_pending_job_gets_retention_ttl_and_can_release_capacity(): + store = InMemoryGraphExtractJobStore(max_jobs=1) + job = store.create(_payload()) + + cancelled = store.cancel(job.job_id) + + assert cancelled.status == GraphExtractJobStatus.CANCELLED + assert cancelled.expires_at is not None + cancelled.expires_at = cancelled.finished_at - timedelta(seconds=1) + store.cleanup() + replacement = store.create(_payload()) + assert replacement.job_id != job.job_id + assert len(store.list_jobs()) == 1 + + +def test_unknown_and_expired_jobs_return_explicit_semantics(): + store = InMemoryGraphExtractJobStore(result_ttl_seconds=0) + client = _client(Mock(), store, run_jobs_inline=False) + + created = client.post("/graph/extract/jobs", json=_payload()).json() + store.expire_jobs() + unknown_response = client.get("/graph/extract/jobs/missing") + expired_result = client.get(f"/graph/extract/jobs/{created['job_id']}/result") + + assert unknown_response.status_code == status.HTTP_404_NOT_FOUND + assert unknown_response.json()["detail"]["code"] == "GRAPH_EXTRACT_JOB_NOT_FOUND" + assert expired_result.status_code == status.HTTP_410_GONE + assert expired_result.json()["detail"]["code"] == "GRAPH_EXTRACT_JOB_EXPIRED" + + +def test_expiring_job_clears_result_payload(): + store = InMemoryGraphExtractJobStore(result_ttl_seconds=0) + job = store.create(_payload()) + store.mark_running(job.job_id) + store.mark_succeeded(job.job_id, _success_response()) + + store.expire_jobs() + + expired_job = store.get(job.job_id) + assert expired_job.status == GraphExtractJobStatus.EXPIRED + assert expired_job.result is None + + +def test_expired_jobs_do_not_count_against_capacity_after_cleanup(): + store = InMemoryGraphExtractJobStore(max_jobs=1, result_ttl_seconds=0) + + first_job = store.create(_payload()) + store.expire_jobs() + second_job = store.create(_payload()) + + assert first_job.status == GraphExtractJobStatus.EXPIRED + assert second_job.job_id != first_job.job_id + assert len(store.list_jobs()) == 1 + + +def test_queue_full_rolls_back_pending_job(monkeypatch): + store = InMemoryGraphExtractJobStore(max_jobs=1) + job = store.create(_payload()) + monkeypatch.setattr(store._queue, "put_nowait", Mock(side_effect=Full)) + + try: + store.submit_job(job.job_id, Mock()) + except ValueError as exc: + assert "queue is full" in str(exc) + else: + raise AssertionError("queue.Full should fail job submission") + + assert store.get(job.job_id) is None + replacement = store.create(_payload()) + assert replacement.job_id != job.job_id + + +def test_running_job_cancellation_does_not_claim_task_was_stopped(): + store = InMemoryGraphExtractJobStore() + job = store.create(_payload()) + store.mark_running(job.job_id) + + cancelled = store.cancel(job.job_id) + + assert cancelled.status == GraphExtractJobStatus.RUNNING + + +def test_delete_running_job_returns_explicit_not_cancellable_error(): + store = InMemoryGraphExtractJobStore() + job = store.create(_payload()) + store.mark_running(job.job_id) + client = _client(Mock(), store, run_jobs_inline=False) + + response = client.delete(f"/graph/extract/jobs/{job.job_id}") + + assert response.status_code == status.HTTP_409_CONFLICT + assert response.json()["detail"]["code"] == "GRAPH_EXTRACT_JOB_NOT_CANCELLABLE" + assert store.get(job.job_id).status == GraphExtractJobStatus.RUNNING + + +def test_zero_ttl_does_not_expire_running_job(): + store = InMemoryGraphExtractJobStore(result_ttl_seconds=0) + job = store.create(_payload()) + store.mark_running(job.job_id) + + store.expire_jobs() + + assert store.get(job.job_id).status == GraphExtractJobStatus.RUNNING + + +def test_concurrent_job_creation_preserves_store_state(): + store = InMemoryGraphExtractJobStore(max_jobs=20) + + with ThreadPoolExecutor(max_workers=5) as executor: + job_ids = list(executor.map(lambda _: store.create(_payload()).job_id, range(10))) + + assert len(set(job_ids)) == 10 + assert len(store.list_jobs()) == 10 diff --git a/hugegraph-llm/src/tests/api/test_graph_extract_routes.py b/hugegraph-llm/src/tests/api/test_graph_extract_routes.py new file mode 100644 index 000000000..27629ecab --- /dev/null +++ b/hugegraph-llm/src/tests/api/test_graph_extract_routes.py @@ -0,0 +1,84 @@ +# 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 unittest.mock import Mock + +from fastapi import APIRouter, FastAPI + +from hugegraph_llm.api.graph_extract_api import graph_extract_http_api +from hugegraph_llm.api.rag_api import rag_http_api + + +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) + + openapi_paths = app.openapi()["paths"] + paths = set(openapi_paths) + assert "/rag" in paths + assert "/text2gremlin" in paths + assert "/config/graph" in paths + assert "/graph/extract" in paths + assert "/graph/extract/jobs" in paths + assert "/graph/import" in paths + assert "/graph/extract-and-import" in paths + import_schema_ref = openapi_paths["/graph/import"]["post"]["responses"]["200"]["content"]["application/json"][ + "schema" + ]["$ref"] + extract_import_schema_ref = openapi_paths["/graph/extract-and-import"]["post"]["responses"]["200"]["content"][ + "application/json" + ]["schema"]["$ref"] + assert import_schema_ref.endswith("/GraphImportResponse") + assert extract_import_schema_ref.endswith("/GraphExtractAndImportResponse") + + +def test_rag_demo_registers_graph_extract_routes_once(monkeypatch): + from hugegraph_llm.demo.rag_demo import app as rag_demo_app + + monkeypatch.setattr(rag_demo_app.prompt, "update_yaml_file", lambda: None) + monkeypatch.setattr(rag_demo_app, "init_rag_ui", lambda: object()) + monkeypatch.setattr(rag_demo_app.gr, "mount_gradio_app", lambda app, *args, **kwargs: app) + + app = rag_demo_app.create_app() + + graph_route_methods = [ + (path, method.upper()) + for path, path_item in app.openapi()["paths"].items() + if path.startswith("/graph/") + for method in path_item + if method.upper() in {"GET", "POST", "DELETE"} + ] + assert len(graph_route_methods) == len(set(graph_route_methods)) + assert ("/graph/extract", "POST") in graph_route_methods + assert ("/graph/extract/jobs", "POST") in graph_route_methods + assert ("/graph/extract/jobs/{job_id}", "GET") in graph_route_methods + assert ("/graph/extract/jobs/{job_id}", "DELETE") in graph_route_methods + assert ("/graph/extract/jobs/{job_id}/result", "GET") in graph_route_methods + assert ("/graph/import", "POST") in graph_route_methods + assert ("/graph/extract-and-import", "POST") in graph_route_methods diff --git a/hugegraph-llm/src/tests/api/test_graph_import_api.py b/hugegraph-llm/src/tests/api/test_graph_import_api.py new file mode 100644 index 000000000..8fb40f40a --- /dev/null +++ b/hugegraph-llm/src/tests/api/test_graph_import_api.py @@ -0,0 +1,744 @@ +# 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 Mock + +import pytest +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_requests import GraphExtractAndImportRequest, GraphImportRequest +from hugegraph_llm.api.models.graph_extract_responses import GraphExtractResponse, GraphImportResponse +from hugegraph_llm.config import huge_settings +from hugegraph_llm.flows import FlowName +from hugegraph_llm.services.graph_extract_service import ( + FlowOutputValidationError, + GraphExtractService, + GraphImportService, + apply_client_config, +) + + +def _payload_data(): + return { + "vertices": [{"label": "person", "properties": {"name": "marko"}}], + "edges": [ + { + "label": "knows", + "outV": "marko", + "outVLabel": "person", + "inV": "vadas", + "inVLabel": "person", + "properties": {}, + } + ], + } + + +def _target_client_config(graph="target_graph"): + return {"graph": graph, "user": "admin", "pwd": "secret", "gs": "space_a"} + + +def _import_payload(**overrides): + payload = { + "schema": { + "propertykeys": [ + {"name": "name", "data_type": "TEXT", "cardinality": "SINGLE"}, + ], + "vertexlabels": [{"name": "person", "properties": ["name"]}], + "edgelabels": [{"name": "knows", "source_label": "person", "target_label": "person"}], + }, + "data": _payload_data(), + "write_to_graph": True, + "client_config": _target_client_config(), + } + payload.update(overrides) + return payload + + +def _extract_payload(**overrides): + payload = { + "texts": ["marko knows vadas"], + "schema": { + "vertexlabels": [{"name": "person", "properties": ["name"]}], + "edgelabels": [{"name": "knows", "source_label": "person", "target_label": "person"}], + }, + "example_prompt": "extract graph", + "write_to_graph": True, + } + payload.update(overrides) + return payload + + +def _client(extract_service=None, import_service=None): + router = APIRouter() + graph_extract_http_api(router, service=extract_service, import_service=import_service) + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def _import_response(updated_embeddings=False): + return GraphImportResponse( + status="succeeded", + vertex_count=1, + edge_count=1, + updated_embeddings=updated_embeddings, + warnings=[], + meta={"duration_ms": 2}, + ) + + +def _flow_import_result( + vertices_created=1, + edges_created=1, + triples_created=0, + vertices_skipped=0, + edges_skipped=0, + triples_skipped=0, + errors=None, +): + return json.dumps( + { + "import_result": { + "vertices_attempted": vertices_created + vertices_skipped, + "vertices_created": vertices_created, + "vertices_skipped": vertices_skipped, + "edges_attempted": edges_created + edges_skipped, + "edges_created": edges_created, + "edges_skipped": edges_skipped, + "triples_attempted": triples_created + triples_skipped, + "triples_created": triples_created, + "triples_skipped": triples_skipped, + "errors": errors or [], + } + } + ) + + +def _partial_import_response(): + return GraphImportResponse( + status="partial", + vertex_count=1, + edge_count=0, + updated_embeddings=False, + warnings=["missing primary key"], + meta={"duration_ms": 2, "import_result": {"vertices_created": 1, "edges_skipped": 1}}, + ) + + +def test_post_graph_import_calls_import_service(): + import_service = Mock() + import_service.import_graph.return_value = _import_response() + client = _client(import_service=import_service) + + response = client.post("/graph/import", json=_import_payload()) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["vertex_count"] == 1 + assert response.json()["updated_embeddings"] is False + import_service.import_graph.assert_called_once() + + +def test_post_graph_import_requires_write_confirmation_before_writing(): + import_service = Mock() + client = _client(import_service=import_service) + + response = client.post("/graph/import", json=_import_payload(write_to_graph=False)) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert response.json()["detail"]["code"] == "GRAPH_IMPORT_CONFIRMATION_REQUIRED" + import_service.import_graph.assert_not_called() + + +def test_post_graph_import_validation_error_uses_import_code(): + import_service = Mock() + client = _client(import_service=import_service) + + response = client.post( + "/graph/import", + json=_import_payload(data={"vertices": [{"label": "person", "properties": None}]}), + ) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + assert response.json()["detail"]["code"] == "GRAPH_IMPORT_VALIDATION_ERROR" + assert response.json()["detail"]["phase"] == "import" + import_service.import_graph.assert_not_called() + + +def test_graph_import_request_rejects_empty_graph_data(): + with pytest.raises(ValueError, match="vertex or edge"): + GraphImportRequest(schema={"vertices": [{"label": "person"}]}, data={"vertices": [], "edges": []}) + + +def test_graph_import_request_rejects_triples_only_graph_data(): + with pytest.raises(ValueError, match="triples import is not supported"): + GraphImportRequest( + schema={ + "vertexlabels": [{"name": "person", "properties": ["name"]}], + "edgelabels": [], + }, + data={"triples": [{"start": "marko", "type": "knows", "end": "vadas"}]}, + write_to_graph=True, + ) + + +def test_graph_import_request_rejects_malformed_property_graph_data(): + invalid_payloads = [ + {"vertices": "not-a-list"}, + {"vertices": [{"properties": {"name": "marko"}}]}, + {"vertices": [{"label": "person", "properties": "not-an-object"}]}, + {"edges": "not-a-list"}, + {"edges": [{"label": "knows", "outV": "marko", "inV": "vadas", "properties": {}}]}, + { + "edges": [ + { + "label": "knows", + "outV": "marko", + "outVLabel": "person", + "inV": "vadas", + "inVLabel": "person", + "properties": "not-an-object", + } + ] + }, + ] + + for data in invalid_payloads: + with pytest.raises(ValueError): + GraphImportRequest(**_import_payload(data=data)) + + +def test_graph_import_request_rejects_edge_endpoint_label_mismatch(): + data = _payload_data() + data["edges"][0]["outVLabel"] = "movie" + + with pytest.raises(ValueError, match="outVLabel must match schema source_label"): + GraphImportRequest(**_import_payload(data=data)) + + +def test_graph_import_request_rejects_inline_schema_unknown_property(): + data = _payload_data() + data["vertices"][0]["properties"]["secret"] = "raw user data" + + with pytest.raises(ValueError, match="vertices\\[0\\].properties.secret"): + GraphImportRequest(**_import_payload(data=data)) + + +def test_graph_import_request_rejects_inline_schema_invalid_property_type(): + data = _payload_data() + data["vertices"][0]["properties"]["name"] = 123 + + with pytest.raises(ValueError, match="vertices\\[0\\].properties.name must match schema property type"): + GraphImportRequest(**_import_payload(data=data)) + + +def test_graph_import_request_rejects_bool_for_integer_property_and_accepts_integer_double(): + payload = _import_payload( + schema={ + "propertykeys": [ + {"name": "name", "data_type": "TEXT", "cardinality": "SINGLE"}, + {"name": "age", "data_type": "INT", "cardinality": "SINGLE"}, + {"name": "score", "data_type": "DOUBLE", "cardinality": "SINGLE"}, + ], + "vertexlabels": [{"name": "person", "properties": ["name", "age", "score"]}], + "edgelabels": [], + }, + data={"vertices": [{"label": "person", "properties": {"name": "marko", "age": 29, "score": 1}}]}, + ) + + request = GraphImportRequest(**payload) + assert request.data["vertices"][0]["properties"]["score"] == 1 + + payload["data"]["vertices"][0]["properties"]["age"] = True + with pytest.raises(ValueError, match="vertices\\[0\\].properties.age must match schema property type"): + GraphImportRequest(**payload) + + +def test_graph_import_request_rejects_inline_schema_unknown_edge_property(): + data = _payload_data() + data["edges"][0]["properties"]["secret"] = "raw edge data" + + with pytest.raises(ValueError, match="edges\\[0\\].properties.secret"): + GraphImportRequest(**_import_payload(data=data)) + + +def test_graph_import_request_requires_explicit_target_graph_for_inline_schema_write(): + with pytest.raises(ValueError, match="client_config.graph"): + GraphImportRequest(**_import_payload(client_config=None)) + + with pytest.raises(ValueError, match="client_config.graph"): + GraphExtractAndImportRequest(**_extract_payload(write_to_graph=True)) + + +def test_extract_and_import_requires_write_confirmation_before_writing(): + import_service = Mock() + client = _client(import_service=import_service) + + response = client.post("/graph/extract-and-import", json=_extract_payload(write_to_graph=False)) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert response.json()["detail"]["code"] == "GRAPH_IMPORT_CONFIRMATION_REQUIRED" + import_service.import_graph.assert_not_called() + + +def test_confirmed_extract_and_import_runs_extraction_before_import(): + extract_service = Mock() + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={"vertex_count": 1, "edge_count": 1}, + ) + import_service = Mock() + import_service.import_graph.return_value = _import_response() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post("/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config())) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["status"] == "succeeded" + assert response.json()["import_result"]["status"] == "succeeded" + extract_service.extract_sync.assert_called_once() + import_service.import_graph.assert_called_once() + + +def test_extract_and_import_top_status_tracks_import_result(): + extract_service = Mock() + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={}, + ) + import_service = Mock() + import_service.import_graph.return_value = _partial_import_response() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post("/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config())) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["status"] == "partial" + assert response.json()["import_result"]["status"] == "partial" + + +def test_extract_and_import_extract_phase_runtime_error_is_sanitized(): + extract_service = Mock() + extract_service.extract_sync.side_effect = RuntimeError("raw text provider timeout secret") + import_service = Mock() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post("/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config())) + + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["detail"] == { + "code": "GRAPH_EXTRACT_FAILED", + "message": "Graph extraction failed during execution", + "phase": "extract", + } + assert "raw text provider timeout secret" not in str(response.json()) + import_service.import_graph.assert_not_called() + + +def test_extract_and_import_import_phase_runtime_error_is_sanitized(): + extract_service = Mock() + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={}, + ) + import_service = Mock() + import_service.import_graph.side_effect = RuntimeError("hugegraph password secret") + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post("/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config())) + + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["detail"] == { + "code": "GRAPH_IMPORT_FAILED", + "message": "Graph import failed during execution", + "phase": "import", + } + assert "hugegraph password secret" not in str(response.json()) + + +def test_extract_and_import_flow_output_errors_keep_phase_specific_codes(): + extract_service = Mock() + extract_service.extract_sync.side_effect = FlowOutputValidationError("raw llm output") + import_service = Mock() + client = _client(extract_service=extract_service, import_service=import_service) + + extract_error = client.post( + "/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config()) + ) + + assert extract_error.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert extract_error.json()["detail"] == { + "code": "GRAPH_EXTRACT_INVALID_FLOW_OUTPUT", + "message": "Graph extraction flow output is invalid", + "phase": "extract", + } + import_service.import_graph.assert_not_called() + + extract_service.extract_sync.side_effect = None + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={}, + ) + import_service.import_graph.side_effect = FlowOutputValidationError("raw import output") + + import_error = client.post( + "/graph/extract-and-import", json=_extract_payload(client_config=_target_client_config()) + ) + + assert import_error.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert import_error.json()["detail"] == { + "code": "GRAPH_IMPORT_INVALID_FLOW_OUTPUT", + "message": "Graph import flow output is invalid", + "phase": "import", + } + + +def test_extract_and_import_allows_inline_schema_with_request_graph_config(): + request = GraphExtractAndImportRequest(**_extract_payload(client_config=_target_client_config())) + + assert request.client_config.graph == "target_graph" + + +def test_extract_and_import_passes_inline_schema_request_graph_config_to_import(): + extract_service = Mock() + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={}, + ) + import_service = Mock() + import_service.import_graph.return_value = _import_response() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post( + "/graph/extract-and-import", + json=_extract_payload(client_config=_target_client_config()), + ) + + assert response.status_code == status.HTTP_200_OK + import_request = import_service.import_graph.call_args.args[0] + assert import_request.client_config.graph == "target_graph" + assert import_request.client_config.user == "admin" + assert import_request.client_config.pwd == "secret" + assert import_request.client_config.gs == "space_a" + + +def test_extract_and_import_inline_schema_keeps_client_config_out_of_extract_flow(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = '{"vertices":[],"edges":[]}' + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + request = GraphExtractAndImportRequest(**_extract_payload(client_config=_target_client_config())) + + GraphExtractService().extract_sync(request) + + assert scheduler.schedule_flow.call_args.kwargs["client_config"] is None + + +def test_extract_and_import_rejects_triples_extraction_at_request_boundary(): + extract_service = Mock() + import_service = Mock() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post( + "/graph/extract-and-import", + json=_extract_payload(extract_type="triples", client_config=_target_client_config()), + ) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY + extract_service.extract_sync.assert_not_called() + import_service.import_graph.assert_not_called() + + +def test_graph_import_service_uses_import_flow_and_updates_embeddings_only_when_requested(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.side_effect = [_flow_import_result(), "{}"] + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + response = GraphImportService().import_graph( + GraphImportRequest(**_import_payload(options={"update_vid_embeddings": True})) + ) + + assert response.updated_embeddings is True + assert scheduler.schedule_flow.call_args_list[0].args[0] == FlowName.IMPORT_GRAPH_DATA + assert scheduler.schedule_flow.call_args_list[1].args[0] == FlowName.UPDATE_VID_EMBEDDINGS + assert scheduler.schedule_flow.call_args_list[1].kwargs["graph_config"]["graph"] == "target_graph" + + +def test_graph_import_service_requires_write_confirmation(): + scheduler = Mock() + + with pytest.raises(ValueError, match="write_to_graph must be True"): + GraphImportService(scheduler=scheduler).import_graph( + GraphImportRequest(**_import_payload(write_to_graph=False, client_config=None)) + ) + + scheduler.schedule_flow.assert_not_called() + + +def test_graph_import_service_keeps_import_result_when_embedding_update_fails(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.side_effect = [ + _flow_import_result(), + RuntimeError("embed failed at http://internal.example/token=secret"), + ] + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + response = GraphImportService().import_graph( + GraphImportRequest(**_import_payload(options={"update_vid_embeddings": True})) + ) + + assert response.status == "partial" + assert response.vertex_count == 1 + assert response.edge_count == 1 + assert response.updated_embeddings is False + assert response.warnings == ["update_vid_embeddings failed"] + assert "internal.example" not in " ".join(response.warnings) + assert "secret" not in " ".join(response.warnings) + assert scheduler.schedule_flow.call_args_list[0].args[0] == FlowName.IMPORT_GRAPH_DATA + assert scheduler.schedule_flow.call_args_list[1].args[0] == FlowName.UPDATE_VID_EMBEDDINGS + + +def test_graph_import_service_does_not_update_embeddings_by_default(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = _flow_import_result() + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + response = GraphImportService().import_graph(GraphImportRequest(**_import_payload())) + + assert response.updated_embeddings is False + assert len(scheduler.schedule_flow.call_args_list) == 1 + assert scheduler.schedule_flow.call_args.args[0] == FlowName.IMPORT_GRAPH_DATA + + +def test_graph_import_service_uses_request_graph_config_without_mutating_global_config(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = _flow_import_result(edges_created=0) + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + monkeypatch.setattr(huge_settings, "graph_url", "127.0.0.1:8080") + monkeypatch.setattr(huge_settings, "graph_name", "before-graph") + monkeypatch.setattr(huge_settings, "graph_user", "before-user") + monkeypatch.setattr(huge_settings, "graph_pwd", "before-pwd") + monkeypatch.setattr(huge_settings, "graph_space", "before-space") + + response = GraphImportService().import_graph( + GraphImportRequest( + **_import_payload( + client_config={ + "graph": "hugegraph", + "user": "admin", + "pwd": "secret", + "gs": "space_a", + } + ) + ) + ) + + assert scheduler.schedule_flow.call_args.kwargs["graph_config"] == { + "graph": "hugegraph", + "user": "admin", + "pwd": "secret", + "gs": "space_a", + } + assert response.meta["client_config"]["pwd"] == "***" + assert "secret" not in str(response.model_dump()) + assert huge_settings.graph_url == "127.0.0.1:8080" + assert huge_settings.graph_name == "before-graph" + assert huge_settings.graph_user == "before-user" + assert huge_settings.graph_pwd == "before-pwd" + assert huge_settings.graph_space == "before-space" + + +def test_graph_import_service_aligns_graph_name_schema_with_write_target(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = _flow_import_result(vertices_created=0, edges_created=0) + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + GraphImportService().import_graph( + GraphImportRequest( + schema="tenant_graph", + data=_payload_data(), + write_to_graph=True, + client_config={"user": "admin"}, + ) + ) + + assert scheduler.schedule_flow.call_args.kwargs["graph_config"] == { + "graph": "tenant_graph", + "user": "admin", + } + + +def test_apply_client_config_accepts_dict_without_mutating_input(): + config = {"graph": "tenant_graph", "user": "admin", "pwd": None} + + result = apply_client_config(config, schema="tenant_graph", align_graph_with_schema=True) + + assert result == {"graph": "tenant_graph", "user": "admin"} + assert config == {"graph": "tenant_graph", "user": "admin", "pwd": None} + + +def test_graph_import_service_rejects_schema_graph_and_target_graph_mismatch(monkeypatch): + scheduler = Mock() + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + with pytest.raises(ValueError, match="schema graph name"): + GraphImportRequest( + schema="schema_graph", + data=_payload_data(), + write_to_graph=True, + client_config={"graph": "target_graph"}, + ) + + scheduler.schedule_flow.assert_not_called() + + +def test_graph_import_service_passes_graph_config_to_vid_embedding_update(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.side_effect = [ + _flow_import_result(vertices_created=0, edges_created=0), + "{}", + ] + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + GraphImportService().import_graph( + GraphImportRequest( + **_import_payload( + client_config={ + "graph": "tenant_graph", + }, + options={"update_vid_embeddings": True}, + ) + ) + ) + + assert scheduler.schedule_flow.call_args_list[1].kwargs["graph_config"] == { + "graph": "tenant_graph", + } + + +def test_graph_import_response_counts_actual_import_result_not_request_size(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = ( + '{"import_result":{"vertices_attempted":1,"vertices_created":0,"vertices_skipped":1,' + '"edges_attempted":1,"edges_created":1,"edges_skipped":0,' + '"triples_attempted":0,"triples_created":0,"triples_skipped":0,' + '"errors":[{"kind":"vertex","index":0,"label":"person","key":"name","reason":"missing_primary_key"}]}}' + ) + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + response = GraphImportService().import_graph(GraphImportRequest(**_import_payload())) + + assert response.status == "partial" + assert response.vertex_count == 0 + assert response.edge_count == 1 + assert response.meta["import_result"]["vertices_skipped"] == 1 + assert response.meta["import_result"]["errors"] == [ + {"kind": "vertex", "index": 0, "label": "person", "key": "name", "reason": "missing_primary_key"} + ] + assert response.warnings == ["vertex import error index=0 label=person key=name reason=missing_primary_key"] + + +def test_graph_import_response_marks_all_skipped_result_as_failed(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = ( + '{"import_result":{"vertices_attempted":0,"vertices_created":0,"vertices_skipped":1,' + '"edges_attempted":0,"edges_created":0,"edges_skipped":0,' + '"triples_attempted":0,"triples_created":0,"triples_skipped":0,' + '"errors":["vertex creation failed for Tom Hanks with token=secret"]}}' + ) + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + response = GraphImportService().import_graph(GraphImportRequest(**_import_payload())) + + assert response.status == "failed" + assert response.vertex_count == 0 + assert response.meta["import_result"]["errors"] == [{"kind": "import", "reason": "import_error"}] + assert response.warnings == ["import error reason=import_error"] + assert "Tom Hanks" not in str(response.warnings) + assert "secret" not in str(response.meta["import_result"]["errors"]) + + +@pytest.mark.parametrize("raw_result", ["{}", '{"import_result":{}}']) +def test_graph_import_service_rejects_missing_or_malformed_import_result(monkeypatch, raw_result): + scheduler = Mock() + scheduler.schedule_flow.return_value = raw_result + monkeypatch.setattr( + "hugegraph_llm.services.graph_extract_service.SchedulerSingleton.get_instance", + lambda: scheduler, + ) + + with pytest.raises(FlowOutputValidationError, match="import_result"): + GraphImportService().import_graph(GraphImportRequest(**_import_payload())) + + +def test_extract_endpoint_never_invokes_import_service(): + extract_service = Mock() + extract_service.extract_sync.return_value = GraphExtractResponse( + status="succeeded", + result=_payload_data(), + warnings=[], + meta={"vertex_count": 1, "edge_count": 1}, + ) + import_service = Mock() + client = _client(extract_service=extract_service, import_service=import_service) + + response = client.post("/graph/extract", json=_extract_payload()) + + assert response.status_code == status.HTTP_200_OK + import_service.import_graph.assert_not_called() diff --git a/hugegraph-llm/src/tests/document/test_graph_extract_configurable_split.py b/hugegraph-llm/src/tests/document/test_graph_extract_configurable_split.py index 4e5078bb9..969ea6496 100644 --- a/hugegraph-llm/src/tests/document/test_graph_extract_configurable_split.py +++ b/hugegraph-llm/src/tests/document/test_graph_extract_configurable_split.py @@ -45,6 +45,7 @@ def to_json(self): "chunks": ["chunk one", "chunk two"], "vertices": [{"id": "person:alice"}], "edges": [], + "max_parallel_chunks": 2, } @@ -205,6 +206,7 @@ def test_graph_extract_post_deal_logs_chunk_count(monkeypatch): result_data = json.loads(result) assert result_data["vertices"] == [{"id": "person:alice"}] + assert result_data["max_parallel_chunks"] == 2 assert any(message == "Graph extraction chunk_count: %s" and args == (2,) for message, args in log_calls) diff --git a/hugegraph-llm/src/tests/flows/test_graph_extract_flow.py b/hugegraph-llm/src/tests/flows/test_graph_extract_flow.py new file mode 100644 index 000000000..382108e1f --- /dev/null +++ b/hugegraph-llm/src/tests/flows/test_graph_extract_flow.py @@ -0,0 +1,140 @@ +# 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 hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + +SCHEMA = '{"vertices": [], "edges": []}' +TEXTS = ["Alice knows Bob."] +EXAMPLE_PROMPT = "" +EXTRACT_TYPE = "property_graph" + + +def test_prepare_writes_requested_split_type(): + flow = GraphExtractFlow() + prepared_input = WkFlowInput() + + flow.prepare( + prepared_input, + SCHEMA, + TEXTS, + EXAMPLE_PROMPT, + EXTRACT_TYPE, + split_type="paragraph", + ) + + assert prepared_input.split_type == "paragraph" + + +def test_prepare_writes_content_type_and_parallel_chunks(): + flow = GraphExtractFlow() + prepared_input = WkFlowInput() + + flow.prepare( + prepared_input, + SCHEMA, + ["chunk-a", "chunk-b"], + EXAMPLE_PROMPT, + EXTRACT_TYPE, + content_type="chunks", + max_parallel_chunks=4, + ) + + assert prepared_input.texts == ["chunk-a", "chunk-b"] + assert prepared_input.content_type == "chunks" + assert prepared_input.max_parallel_chunks == 4 + + +def test_prepare_rejects_chunks_with_non_document_split_type(): + flow = GraphExtractFlow() + + try: + flow.prepare( + WkFlowInput(), + SCHEMA, + ["chunk-a"], + EXAMPLE_PROMPT, + EXTRACT_TYPE, + split_type="paragraph", + content_type="chunks", + ) + except ValueError as exc: + assert "split_type must be document when content_type is chunks" in str(exc) + else: + raise AssertionError("chunks content must reject non-document split_type") + + +def test_chunk_split_node_uses_pre_split_chunks_without_splitting(): + node = ChunkSplitNode() + node.wk_input = WkFlowInput() + node.wk_input.texts = ["chunk-a\n\nchunk-b", "chunk-c"] + node.wk_input.language = "en" + node.wk_input.split_type = "document" + node.wk_input.content_type = "chunks" + node.context = WkFlowState() + + status = node.node_init() + result = node.operator_schedule({}) + + assert not status.isErr() + assert node.chunk_split_op is None + assert result["chunks"] == ["chunk-a\n\nchunk-b", "chunk-c"] + + +def test_build_flow_writes_requested_split_type_to_workflow_input(): + flow = GraphExtractFlow() + + pipeline = flow.build_flow( + SCHEMA, + TEXTS, + EXAMPLE_PROMPT, + EXTRACT_TYPE, + split_type="paragraph", + ) + + wkflow_input = pipeline.getGParamWithNoEmpty("wkflow_input") + assert wkflow_input.split_type == "paragraph" + + +def test_build_flow_defaults_to_document_split_type_for_existing_callers(): + flow = GraphExtractFlow() + + pipeline = flow.build_flow(SCHEMA, TEXTS, EXAMPLE_PROMPT, EXTRACT_TYPE) + + wkflow_input = pipeline.getGParamWithNoEmpty("wkflow_input") + assert wkflow_input.split_type == "document" + + +def test_workflow_state_setup_clears_graph_extract_result_fields(): + state = WkFlowState() + state.vertices = [{"id": "old"}] + state.edges = [{"id": "old-edge"}] + state.triples = [("old", "rel", "value")] + state.chunks = ["old chunk"] + state.call_count = 10 + state.max_parallel_chunks = 2 + + status = state.setup() + + assert not status.isErr() + assert state.vertices is None + assert state.edges is None + assert state.triples is None + assert state.chunks is None + assert state.call_count is None + assert state.max_parallel_chunks is None diff --git a/hugegraph-llm/src/tests/models/llms/test_init_llm.py b/hugegraph-llm/src/tests/models/llms/test_init_llm.py new file mode 100644 index 000000000..6ca748ae0 --- /dev/null +++ b/hugegraph-llm/src/tests/models/llms/test_init_llm.py @@ -0,0 +1,189 @@ +# 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 types import SimpleNamespace +from unittest.mock import patch + +from hugegraph_llm.models.llms.init_llm import get_extract_llm + + +def _openai_config(**overrides): + config = { + "chat_llm_type": "openai", + "extract_llm_type": "openai", + "openai_chat_api_key": "chat-key", + "openai_chat_api_base": "https://chat.example/v1", + "openai_chat_language_model": "chat-model", + "openai_chat_tokens": 4096, + "openai_extract_api_key": None, + "openai_extract_api_base": "https://api.openai.com/v1", + "openai_extract_language_model": "gpt-4.1-mini", + "openai_extract_tokens": 256, + } + config.update(overrides) + return SimpleNamespace(**config) + + +def _litellm_config(**overrides): + config = { + "chat_llm_type": "litellm", + "extract_llm_type": "litellm", + "litellm_chat_api_key": "chat-key", + "litellm_chat_api_base": "https://chat.example/v1", + "litellm_chat_language_model": "chat-model", + "litellm_chat_tokens": 4096, + "litellm_extract_api_key": None, + "litellm_extract_api_base": None, + "litellm_extract_language_model": "openai/gpt-4.1-mini", + "litellm_extract_tokens": 256, + } + config.update(overrides) + return SimpleNamespace(**config) + + +def _ollama_config(**overrides): + config = { + "chat_llm_type": "ollama/local", + "extract_llm_type": "ollama/local", + "ollama_chat_host": "chat-host", + "ollama_chat_port": 11435, + "ollama_chat_language_model": "chat-model", + "ollama_extract_host": "127.0.0.1", + "ollama_extract_port": 11434, + "ollama_extract_language_model": None, + } + config.update(overrides) + return SimpleNamespace(**config) + + +def test_get_extract_llm_falls_back_to_openai_chat_config_when_extract_key_is_missing(): + config = _openai_config() + + with patch("hugegraph_llm.models.llms.init_llm.OpenAIClient") as openai_client: + get_extract_llm(config) + + openai_client.assert_called_once_with( + api_key="chat-key", + api_base="https://chat.example/v1", + model_name="chat-model", + max_tokens=4096, + ) + + +def test_get_extract_llm_falls_back_to_openai_chat_config_when_only_generic_key_exists(): + config = _openai_config( + openai_chat_api_key="shared-key", + openai_extract_api_key="shared-key", + ) + + with patch("hugegraph_llm.models.llms.init_llm.OpenAIClient") as openai_client: + get_extract_llm(config) + + openai_client.assert_called_once_with( + api_key="shared-key", + api_base="https://chat.example/v1", + model_name="chat-model", + max_tokens=4096, + ) + + +def test_get_extract_llm_prefers_explicit_openai_extract_config(): + config = _openai_config( + openai_extract_api_key="extract-key", + openai_extract_api_base="https://extract.example/v1", + openai_extract_language_model="extract-model", + openai_extract_tokens=8192, + ) + + with patch("hugegraph_llm.models.llms.init_llm.OpenAIClient") as openai_client: + get_extract_llm(config) + + openai_client.assert_called_once_with( + api_key="extract-key", + api_base="https://extract.example/v1", + model_name="extract-model", + max_tokens=8192, + ) + + +def test_get_extract_llm_falls_back_to_litellm_chat_config_when_extract_key_is_missing(): + config = _litellm_config() + + with patch("hugegraph_llm.models.llms.init_llm.LiteLLMClient") as litellm_client: + get_extract_llm(config) + + litellm_client.assert_called_once_with( + api_key="chat-key", + api_base="https://chat.example/v1", + model_name="chat-model", + max_tokens=4096, + ) + + +def test_get_extract_llm_falls_back_to_litellm_chat_config_without_proxy_key(): + config = _litellm_config(litellm_chat_api_key=None) + + with patch("hugegraph_llm.models.llms.init_llm.LiteLLMClient") as litellm_client: + get_extract_llm(config) + + litellm_client.assert_called_once_with( + api_key=None, + api_base="https://chat.example/v1", + model_name="chat-model", + max_tokens=4096, + ) + + +def test_get_extract_llm_prefers_explicit_litellm_extract_config(): + config = _litellm_config( + litellm_extract_api_key="extract-key", + litellm_extract_api_base="https://extract.example/v1", + litellm_extract_language_model="extract-model", + litellm_extract_tokens=8192, + ) + + with patch("hugegraph_llm.models.llms.init_llm.LiteLLMClient") as litellm_client: + get_extract_llm(config) + + litellm_client.assert_called_once_with( + api_key="extract-key", + api_base="https://extract.example/v1", + model_name="extract-model", + max_tokens=8192, + ) + + +def test_get_extract_llm_falls_back_to_ollama_chat_config_when_extract_model_is_missing(): + config = _ollama_config() + + with patch("hugegraph_llm.models.llms.init_llm.OllamaClient") as ollama_client: + get_extract_llm(config) + + ollama_client.assert_called_once_with(model="chat-model", host="chat-host", port=11435) + + +def test_get_extract_llm_prefers_explicit_ollama_extract_config(): + config = _ollama_config( + ollama_extract_host="extract-host", + ollama_extract_port=11436, + ollama_extract_language_model="extract-model", + ) + + with patch("hugegraph_llm.models.llms.init_llm.OllamaClient") as ollama_client: + get_extract_llm(config) + + ollama_client.assert_called_once_with(model="extract-model", host="extract-host", port=11436) diff --git a/hugegraph-llm/src/tests/nodes/test_base_node.py b/hugegraph-llm/src/tests/nodes/test_base_node.py new file mode 100644 index 000000000..5daaee892 --- /dev/null +++ b/hugegraph-llm/src/tests/nodes/test_base_node.py @@ -0,0 +1,37 @@ +# 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 hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class RuntimeFailingNode(BaseNode): + def operator_schedule(self, data_json): + raise RuntimeError("llm provider timeout") + + +def test_base_node_converts_unexpected_operator_exception_to_error_status(): + node = RuntimeFailingNode() + node.wk_input = WkFlowInput() + node.context = WkFlowState() + + status = node.run() + + assert status.isErr() + assert "llm provider timeout" in status.getInfo() + assert "RuntimeFailingNode" in status.getInfo() + assert "Traceback" not in status.getInfo() diff --git a/hugegraph-llm/src/tests/nodes/test_extract_node.py b/hugegraph-llm/src/tests/nodes/test_extract_node.py new file mode 100644 index 000000000..269a28ac6 --- /dev/null +++ b/hugegraph-llm/src/tests/nodes/test_extract_node.py @@ -0,0 +1,38 @@ +# 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 unittest.mock import Mock, patch + +from hugegraph_llm.nodes.llm_node.extract_info import ExtractNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +def test_extract_node_uses_extract_llm_config_for_property_graph(): + llm = Mock() + node = ExtractNode() + node.wk_input = WkFlowInput() + node.wk_input.example_prompt = "extract prompt" + node.wk_input.extract_type = "property_graph" + node.wk_input.max_parallel_chunks = 2 + node.context = WkFlowState() + + with patch("hugegraph_llm.nodes.llm_node.extract_info.get_extract_llm", return_value=llm) as get_extract_llm: + status = node.node_init() + + assert not status.isErr() + get_extract_llm.assert_called_once() + assert node.property_graph_extract.llm is llm diff --git a/hugegraph-llm/src/tests/nodes/test_request_graph_config.py b/hugegraph-llm/src/tests/nodes/test_request_graph_config.py new file mode 100644 index 000000000..cea14cda5 --- /dev/null +++ b/hugegraph-llm/src/tests/nodes/test_request_graph_config.py @@ -0,0 +1,183 @@ +# 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 hugegraph_llm.config import huge_settings +from hugegraph_llm.nodes.hugegraph_node.commit_to_hugegraph import Commit2GraphNode +from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.index_node.build_semantic_index import BuildSemanticIndexNode +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.hugegraph_utils import get_hg_client + +GRAPH_CONFIG = { + "url": "127.0.0.1:8080", + "graph": "custom_graph", + "user": "admin", + "pwd": "secret", + "gs": "space_a", +} + + +def test_schema_node_passes_request_graph_config_to_schema_manager(monkeypatch): + captured = {} + + class FakeSchemaManager: + def __init__(self, graph_name, graph_config=None): + captured["graph_name"] = graph_name + captured["graph_config"] = graph_config + + monkeypatch.setattr("hugegraph_llm.nodes.hugegraph_node.schema.SchemaManager", FakeSchemaManager) + node = SchemaNode() + node.wk_input = WkFlowInput() + node.wk_input.schema = "custom_graph" + node.wk_input.graph_config = GRAPH_CONFIG + node.context = WkFlowState() + + status = node.node_init() + + assert not status.isErr() + assert captured == {"graph_name": "custom_graph", "graph_config": GRAPH_CONFIG} + + +def test_commit_node_passes_request_graph_config_to_commit_operator(monkeypatch): + captured = {} + + class FakeCommit2Graph: + def __init__(self, graph_config=None): + captured["graph_config"] = graph_config + + monkeypatch.setattr("hugegraph_llm.nodes.hugegraph_node.commit_to_hugegraph.Commit2Graph", FakeCommit2Graph) + node = Commit2GraphNode() + node.wk_input = WkFlowInput() + node.wk_input.graph_config = GRAPH_CONFIG + node.context = WkFlowState() + + status = node.node_init() + + assert not status.isErr() + assert captured == {"graph_config": GRAPH_CONFIG} + + +def test_fetch_graph_data_node_uses_request_graph_config(monkeypatch): + captured = {} + + class FakeFetchGraphData: + def __init__(self, client): + captured["client"] = client + + monkeypatch.setattr("hugegraph_llm.nodes.hugegraph_node.fetch_graph_data.FetchGraphData", FakeFetchGraphData) + monkeypatch.setattr( + "hugegraph_llm.nodes.hugegraph_node.fetch_graph_data.get_hg_client", + lambda graph_config=None: {"graph_config": graph_config}, + ) + node = FetchGraphDataNode() + node.wk_input = WkFlowInput() + node.wk_input.graph_config = GRAPH_CONFIG + node.context = WkFlowState() + + status = node.node_init() + + assert not status.isErr() + assert captured == {"client": {"graph_config": GRAPH_CONFIG}} + + +def test_build_semantic_index_node_uses_request_graph_config(monkeypatch): + captured = {} + + class FakeEmbeddings: + def get_embedding(self): + return "embedding" + + class FakeBuildSemanticIndex: + def __init__(self, embedding, vector_index, graph_config=None): + captured["embedding"] = embedding + captured["vector_index"] = vector_index + captured["graph_config"] = graph_config + + monkeypatch.setattr("hugegraph_llm.nodes.index_node.build_semantic_index.Embeddings", FakeEmbeddings) + monkeypatch.setattr("hugegraph_llm.utils.vector_index_utils.get_vector_index_class", lambda _: "vector-index") + monkeypatch.setattr( + "hugegraph_llm.nodes.index_node.build_semantic_index.BuildSemanticIndex", + FakeBuildSemanticIndex, + ) + node = BuildSemanticIndexNode() + node.wk_input = WkFlowInput() + node.wk_input.graph_config = GRAPH_CONFIG + node.context = WkFlowState() + + status = node.node_init() + + assert not status.isErr() + assert captured == { + "embedding": "embedding", + "vector_index": "vector-index", + "graph_config": GRAPH_CONFIG, + } + + +def test_schema_manager_connection_falls_back_for_missing_optional_fields(monkeypatch): + captured = {} + monkeypatch.setattr(huge_settings, "graph_url", "127.0.0.1:8080") + monkeypatch.setattr(huge_settings, "graph_user", "admin") + monkeypatch.setattr(huge_settings, "graph_pwd", "global-secret") + monkeypatch.setattr(huge_settings, "graph_space", "global-space") + + class FakeHugeClient: + def __init__(self, **kwargs): + captured.update(kwargs) + + def schema(self): + return "schema-client" + + monkeypatch.setattr("hugegraph_llm.operators.hugegraph_op.schema_manager.PyHugeClient", FakeHugeClient) + + manager = SchemaManager("custom_graph", connection={"url": "10.0.0.1:8080", "user": None}) + + assert manager.schema == "schema-client" + assert captured == { + "url": "10.0.0.1:8080", + "graph": "custom_graph", + "user": "admin", + "pwd": "global-secret", + "graphspace": "global-space", + } + + +def test_get_hg_client_preserves_explicit_empty_graphspace(monkeypatch): + captured = {} + monkeypatch.setattr(huge_settings, "graph_url", "127.0.0.1:8080") + monkeypatch.setattr(huge_settings, "graph_name", "global-graph") + monkeypatch.setattr(huge_settings, "graph_user", "admin") + monkeypatch.setattr(huge_settings, "graph_pwd", "secret") + monkeypatch.setattr(huge_settings, "graph_space", "global-space") + + class FakeHugeClient: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr("hugegraph_llm.utils.hugegraph_utils.PyHugeClient", FakeHugeClient) + + get_hg_client({"graph": "custom_graph", "gs": ""}) + + assert captured == { + "url": "127.0.0.1:8080", + "graph": "custom_graph", + "user": "admin", + "pwd": "secret", + "graphspace": "", + } diff --git a/hugegraph-llm/src/tests/operators/document_op/test_chunk_split.py b/hugegraph-llm/src/tests/operators/document_op/test_chunk_split.py index e44a10125..825f81d29 100644 --- a/hugegraph-llm/src/tests/operators/document_op/test_chunk_split.py +++ b/hugegraph-llm/src/tests/operators/document_op/test_chunk_split.py @@ -99,6 +99,15 @@ def test_run_paragraph_split(self): self.assertIn("Second paragraph", all_text) self.assertIn("Third paragraph", all_text) + def test_run_paragraph_split_keeps_short_paragraph_boundaries(self): + """Test paragraph split keeps explicit short paragraph boundaries.""" + text_with_short_paragraphs = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph." + chunk_split = ChunkSplit(text_with_short_paragraphs, split_type="paragraph", language="en") + + result = chunk_split.run({}) + + self.assertEqual(result["chunks"], ["First paragraph.", "Second paragraph.", "Third paragraph."]) + def test_run_sentence_split(self): """Test running sentence split.""" # Use a text with more distinct sentences to ensure splitting diff --git a/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph.py b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph.py index 6876b6e28..b6d65776c 100644 --- a/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph.py +++ b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph.py @@ -23,6 +23,7 @@ from pyhugegraph.utils.exceptions import CreateError, NotFoundError from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph +from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract pytestmark = [pytest.mark.unit] @@ -147,17 +148,34 @@ def test_run_without_schema(self, mock_schema_free_mode): mock_schema_free_mode.return_value = None # Create input data - data = {"vertices": self.vertices, "edges": self.edges, "triples": []} + data = {"triples": [["Tom Hanks", "acted_in", "Forrest Gump"]]} # Run the method result = self.commit2graph.run(data) # Verify that schema_free_mode was called - mock_schema_free_mode.assert_called_once_with([]) + mock_schema_free_mode.assert_called_once_with([["Tom Hanks", "acted_in", "Forrest Gump"]]) # Verify the results self.assertEqual(result, data) + def test_run_without_schema_rejects_property_graph_input(self): + """Test schema-free mode rejects vertices/edges to avoid silent data loss.""" + with self.assertRaisesRegex(ValueError, "Schema-free mode only supports triples"): + self.commit2graph.run({"vertices": self.vertices, "edges": self.edges, "triples": []}) + + def test_run_with_schema_rejects_triples_input(self): + """Test schema mode rejects triples to avoid silently dropping them.""" + with self.assertRaisesRegex(ValueError, "Triples input is not supported"): + self.commit2graph.run( + { + "schema": self.schema, + "vertices": self.vertices, + "edges": self.edges, + "triples": [["Tom Hanks", "acted_in", "Forrest Gump"]], + } + ) + @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._check_property_data_type") def test_set_default_property(self, mock_check_property_data_type): """Test _set_default_property method.""" @@ -303,6 +321,328 @@ 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 + def test_run_initializes_customize_string_schema_with_custom_id_strategy(self): + """Test run creates custom string id vertex labels for CUSTOMIZE_STRING schema.""" + schema_mocks = self._setup_schema_mocks() + schema = { + "propertykeys": [{"name": "name", "data_type": "TEXT", "cardinality": "SINGLE"}], + "vertexlabels": [ + { + "name": "person", + "id_strategy": "CUSTOMIZE_STRING", + "primary_keys": ["name"], + "properties": ["name"], + "nullable_keys": [], + } + ], + "edgelabels": [], + } + vertices = [{"id": "marko", "label": "person", "properties": {"name": "marko"}}] + + with ( + patch.object(self.commit2graph, "_create_property"), + patch.object(self.commit2graph, "_handle_graph_creation", return_value=MagicMock(id="marko")), + ): + self.commit2graph.run({"schema": schema, "vertices": vertices, "edges": []}) + + vertex_builder = schema_mocks["vertex_label"].return_value + vertex_builder.useCustomizeStringId.assert_called_once() + vertex_builder.usePrimaryKeyId.assert_not_called() + + @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}, + }, + { + "id": "movie:Forrest Gump", + "label": "movie", + "properties": {"title": "Forrest Gump", "year": 1994}, + }, + ] + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "inV": "movie:Forrest Gump", + } + ] + + self.commit2graph.load_into_graph(vertices, edges, self.schema) + + self.assertEqual(vertices[0]["id"], "1:Tom Hanks") + self.assertEqual(vertices[1]["id"], "2:Forrest Gump") + mock_handle_graph_creation.assert_any_call( + self.commit2graph.client.graph().addEdge, + "acted_in", + "1:Tom Hanks", + "2:Forrest Gump", + {"role": "Forrest Gump"}, + ) + + @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation") + def test_load_into_graph_uses_explicit_customize_string_ids(self, mock_handle_graph_creation): + """Test custom string ids are passed to HugeGraph when schema requires them.""" + mock_handle_graph_creation.side_effect = [ + MagicMock(id="Tom Hanks"), + MagicMock(id="Forrest Gump"), + MagicMock(id="edge_id"), + ] + schema = { + "propertykeys": [ + {"name": "name", "data_type": "TEXT", "cardinality": "SINGLE"}, + {"name": "title", "data_type": "TEXT", "cardinality": "SINGLE"}, + ], + "vertexlabels": [ + { + "id": 7, + "name": "person", + "id_strategy": "CUSTOMIZE_STRING", + "primary_keys": ["name"], + "properties": ["name"], + "nullable_keys": [], + }, + { + "id": 8, + "name": "movie", + "id_strategy": "CUSTOMIZE_STRING", + "primary_keys": ["title"], + "properties": ["title"], + "nullable_keys": [], + }, + ], + "edgelabels": [{"name": "acted_in", "properties": [], "source_label": "person", "target_label": "movie"}], + } + vertices = [ + {"id": "Tom Hanks", "label": "person", "properties": {"name": "Tom Hanks"}}, + {"id": "Forrest Gump", "label": "movie", "properties": {"title": "Forrest Gump"}}, + ] + edges = [ + { + "label": "acted_in", + "properties": {}, + "outV": "Tom Hanks", + "inV": "Forrest Gump", + } + ] + + self.commit2graph.load_into_graph(vertices, edges, schema) + + mock_handle_graph_creation.assert_any_call( + self.commit2graph.client.graph().addVertex, + "person", + {"name": "Tom Hanks"}, + id="Tom Hanks", + ) + mock_handle_graph_creation.assert_any_call( + self.commit2graph.client.graph().addVertex, + "movie", + {"title": "Forrest Gump"}, + id="Forrest Gump", + ) + mock_handle_graph_creation.assert_any_call( + self.commit2graph.client.graph().addEdge, + "acted_in", + "Tom Hanks", + "Forrest Gump", + {}, + ) + + @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation") + def test_load_into_graph_accepts_normalized_extraction_without_item_type(self, mock_handle_graph_creation): + """Test normalized LLM output without type fields can be committed.""" + mock_handle_graph_creation.side_effect = [ + MagicMock(id="1:Tom Hanks"), + MagicMock(id="2:Forrest Gump"), + MagicMock(id="edge_id"), + ] + llm_output = """{ + "vertices": [ + { + "id": "person:Tom Hanks", + "label": "person", + "properties": { + "name": "Tom Hanks", + "age": 67 + } + }, + { + "id": "movie:Forrest Gump", + "label": "movie", + "properties": { + "title": "Forrest Gump", + "year": 1994 + } + } + ], + "edges": [ + { + "label": "acted_in", + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + "properties": { + "role": "Forrest Gump" + } + } + ] + }""" + + items = PropertyGraphExtract(llm=MagicMock())._extract_and_filter_label(self.schema, llm_output) + vertices = [item for item in items if item["type"] == "vertex"] + edges = [item for item in items if item["type"] == "edge"] + self.assertEqual(edges[0]["outV"], "1:Tom Hanks") + self.assertEqual(edges[0]["inV"], "2:Forrest Gump") + + self.commit2graph.load_into_graph(vertices, edges, self.schema) + + mock_handle_graph_creation.assert_any_call( + self.commit2graph.client.graph().addEdge, + "acted_in", + "1:Tom Hanks", + "2:Forrest Gump", + {"role": "Forrest Gump"}, + ) + + @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation") + def test_load_into_graph_with_data_type_validation_failure(self, mock_handle_graph_creation): + """Test load_into_graph method with data type validation failure.""" + # Setup mocks + mock_handle_graph_creation.return_value = MagicMock(id="vertex_id") + + # Create vertices with incorrect data types (strings for INT fields) + vertices = [ + {"label": "person", "properties": {"name": "Tom Hanks", "age": "67"}}, # age should be int, not str + {"label": "movie", "properties": {"title": "Forrest Gump", "year": "1994"}}, # year should be int, not str + ] + + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "inV": "movie:Forrest Gump", + } + ] + + # Call the method - should skip vertices and edges whose batch endpoints failed. + result = self.commit2graph.load_into_graph(vertices, edges, self.schema) + + self.assertEqual(result["vertices_created"], 0) + self.assertEqual(result["vertices_skipped"], 2) + self.assertEqual(result["edges_created"], 0) + self.assertEqual(result["edges_skipped"], 1) + self.assertEqual( + result["errors"][-1], + {"kind": "edge", "index": 0, "reason": "endpoint_vertex_failed", "label": "acted_in", "key": "outV"}, + ) + mock_handle_graph_creation.assert_not_called() + def test_check_property_data_type_success(self): """Test _check_property_data_type method with valid data types.""" # Test TEXT type @@ -310,6 +650,7 @@ def test_check_property_data_type_success(self): # Test INT type self.assertTrue(self.commit2graph._check_property_data_type("INT", "SINGLE", 67)) + self.assertFalse(self.commit2graph._check_property_data_type("INT", "SINGLE", True)) # Test LIST type with valid items self.assertTrue(self.commit2graph._check_property_data_type("TEXT", "LIST", ["hobby1", "hobby2"])) @@ -337,6 +678,7 @@ def test_check_property_data_type_edge_cases(self): # Test FLOAT/DOUBLE type self.assertTrue(self.commit2graph._check_property_data_type("FLOAT", "SINGLE", 3.14)) self.assertTrue(self.commit2graph._check_property_data_type("DOUBLE", "SINGLE", 3.14)) + self.assertTrue(self.commit2graph._check_property_data_type("DOUBLE", "SINGLE", 1)) self.assertFalse(self.commit2graph._check_property_data_type("FLOAT", "SINGLE", "3.14")) # Test DATE type (format: yyyy-MM-dd) diff --git a/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_error_messages.py b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_error_messages.py new file mode 100644 index 000000000..c895aa80f --- /dev/null +++ b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_error_messages.py @@ -0,0 +1,304 @@ +# 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 unittest.mock import MagicMock, patch + +import pytest + +from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph + +pytestmark = [pytest.mark.unit] + + +def _commit_operator(): + mock_client = MagicMock() + mock_client.schema.return_value = MagicMock() + with patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.PyHugeClient", return_value=mock_client): + return Commit2Graph() + + +def _schema(): + return { + "propertykeys": [ + {"name": "name", "data_type": "TEXT", "cardinality": "SINGLE"}, + {"name": "title", "data_type": "TEXT", "cardinality": "SINGLE"}, + {"name": "role", "data_type": "TEXT", "cardinality": "SINGLE"}, + {"name": "age", "data_type": "INT", "cardinality": "SINGLE"}, + {"name": "score", "data_type": "DOUBLE", "cardinality": "SINGLE"}, + ], + "vertexlabels": [ + { + "name": "person", + "properties": ["name", "age", "score"], + "primary_keys": ["name"], + "nullable_keys": ["age", "score"], + "id_strategy": "PRIMARY_KEY", + }, + { + "name": "movie", + "properties": ["title"], + "primary_keys": ["title"], + "nullable_keys": [], + "id_strategy": "PRIMARY_KEY", + }, + ], + "edgelabels": [{"name": "acted_in", "properties": ["role"], "source_label": "person", "target_label": "movie"}], + } + + +def test_import_errors_do_not_include_raw_vertex_or_edge_payloads(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": "Tom Hanks"}}] + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation", return_value=None): + result = commit2graph.load_into_graph(vertices, edges, _schema()) + + assert result["errors"] == [ + {"kind": "vertex", "index": 0, "reason": "create_failed", "label": "person"}, + { + "kind": "edge", + "index": 0, + "reason": "endpoint_vertex_failed", + "label": "acted_in", + "key": "outV", + }, + ] + error_text = str(result["errors"]) + assert "Tom Hanks" not in error_text + assert "Forrest Gump" not in error_text + + +def test_primary_key_import_error_identifies_key_label_and_index_only(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": ""}}] + + result = commit2graph.load_into_graph(vertices, [], _schema()) + + assert result["errors"] == [ + {"kind": "vertex", "index": 0, "reason": "missing_primary_key", "label": "person", "key": "name"} + ] + assert "{'label':" not in str(result["errors"]) + + +def test_schema_free_import_errors_do_not_include_raw_triple_payloads(): + commit2graph = _commit_operator() + triples = [["Alice Sensitive", "knows", "Bob Sensitive"]] + + with patch.object(commit2graph, "_handle_graph_creation", return_value=None): + result = commit2graph.schema_free_mode(triples) + + assert result["errors"] == [{"kind": "triple", "index": 0, "reason": "create_vertices_failed"}] + error_text = str(result["errors"]) + assert "Alice Sensitive" not in error_text + assert "Bob Sensitive" not in error_text + + +def test_import_rejects_edge_endpoint_label_mismatch_before_writing(): + commit2graph = _commit_operator() + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "outVLabel": "movie", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph([], edges, _schema()) + + assert result["edges_created"] == 0 + assert result["edges_skipped"] == 1 + assert result["errors"] == [ + { + "kind": "edge", + "index": 0, + "reason": "source_label_mismatch", + "label": "acted_in", + "key": "outVLabel", + } + ] + mock_handle_graph_creation.assert_not_called() + + +def test_import_rejects_unknown_vertex_property_without_raw_payload(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": "Tom Hanks", "secret": "raw user data"}}] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph(vertices, [], _schema()) + + assert result["vertices_created"] == 0 + assert result["vertices_skipped"] == 1 + assert result["errors"] == [ + {"kind": "vertex", "index": 0, "reason": "unknown_property", "label": "person", "key": "secret"} + ] + assert "Tom Hanks" not in str(result["errors"]) + assert "raw user data" not in str(result["errors"]) + mock_handle_graph_creation.assert_not_called() + + +def test_import_rejects_unknown_edge_property_before_writing(): + commit2graph = _commit_operator() + edges = [ + { + "label": "acted_in", + "properties": {"secret": "raw edge data"}, + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph([], edges, _schema()) + + assert result["edges_created"] == 0 + assert result["edges_skipped"] == 1 + assert result["errors"] == [ + {"kind": "edge", "index": 0, "reason": "unknown_property", "label": "acted_in", "key": "secret"} + ] + assert "raw edge data" not in str(result["errors"]) + mock_handle_graph_creation.assert_not_called() + + +def test_import_rejects_invalid_edge_property_type_before_writing(): + commit2graph = _commit_operator() + edges = [ + { + "label": "acted_in", + "properties": {"role": 123}, + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph([], edges, _schema()) + + assert result["edges_created"] == 0 + assert result["edges_skipped"] == 1 + assert result["errors"] == [ + {"kind": "edge", "index": 0, "reason": "invalid_property_type", "label": "acted_in", "key": "role"} + ] + mock_handle_graph_creation.assert_not_called() + + +def test_import_skips_edge_when_batch_endpoint_vertex_failed(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": "Tom Hanks", "secret": "raw user data"}}] + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph(vertices, edges, _schema()) + + assert result["vertices_created"] == 0 + assert result["vertices_skipped"] == 1 + assert result["edges_created"] == 0 + assert result["edges_skipped"] == 1 + assert result["errors"] == [ + {"kind": "vertex", "index": 0, "reason": "unknown_property", "label": "person", "key": "secret"}, + { + "kind": "edge", + "index": 0, + "reason": "endpoint_vertex_failed", + "label": "acted_in", + "key": "outV", + }, + ] + assert "Tom Hanks" not in str(result["errors"]) + assert "raw user data" not in str(result["errors"]) + mock_handle_graph_creation.assert_not_called() + + +def test_import_keeps_raw_endpoint_fallback_for_existing_vertices(): + commit2graph = _commit_operator() + edges = [ + { + "label": "acted_in", + "properties": {"role": "Forrest Gump"}, + "outV": "person:Tom Hanks", + "outVLabel": "person", + "inV": "movie:Forrest Gump", + "inVLabel": "movie", + } + ] + + with patch.object(commit2graph, "_handle_graph_creation", return_value=MagicMock(id="edge_id")) as mock_create: + result = commit2graph.load_into_graph([], edges, _schema()) + + assert result["edges_created"] == 1 + assert result["edges_skipped"] == 0 + assert result["errors"] == [] + mock_create.assert_called_once_with( + commit2graph.client.graph().addEdge, + "acted_in", + "person:Tom Hanks", + "movie:Forrest Gump", + {"role": "Forrest Gump"}, + ) + + +def test_import_rejects_bool_for_int_property_before_writing(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": "Tom Hanks", "age": True}}] + + with patch.object(commit2graph, "_handle_graph_creation") as mock_handle_graph_creation: + result = commit2graph.load_into_graph(vertices, [], _schema()) + + assert result["vertices_created"] == 0 + assert result["vertices_skipped"] == 1 + assert result["errors"] == [ + {"kind": "vertex", "index": 0, "reason": "invalid_property_type", "label": "person", "key": "age"} + ] + mock_handle_graph_creation.assert_not_called() + + +def test_import_accepts_integer_for_double_property(): + commit2graph = _commit_operator() + vertices = [{"label": "person", "properties": {"name": "Tom Hanks", "score": 1}}] + + with patch.object(commit2graph, "_handle_graph_creation", return_value=MagicMock(id="person:Tom Hanks")): + result = commit2graph.load_into_graph(vertices, [], _schema()) + + assert result["vertices_created"] == 1 + assert result["vertices_skipped"] == 0 + assert result["errors"] == [] diff --git a/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_load_into_graph.py b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_load_into_graph.py index d576f902c..7ea53bcc0 100644 --- a/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_load_into_graph.py +++ b/hugegraph-llm/src/tests/operators/hugegraph_op/test_commit_to_hugegraph_load_into_graph.py @@ -88,7 +88,9 @@ def test_load_into_graph(self, mock_handle_graph_creation, mock_check_property_d "label": "acted_in", "properties": {"role": "Forrest Gump"}, "outV": "person:Tom Hanks", + "outVLabel": "person", "inV": "movie:Forrest Gump", + "inVLabel": "movie", } ] @@ -388,8 +390,8 @@ def test_property_graph_extract_run_preserves_typed_values_for_commit(self, mock self.assertEqual(mock_handle_graph_creation.call_count, 3) @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation") - def test_load_into_graph_raises_explicit_error_when_vertex_creation_fails(self, mock_handle_graph_creation): - """Test failed vertex creation is reported before edge creation.""" + def test_load_into_graph_reports_vertex_creation_failure_and_continues(self, mock_handle_graph_creation): + """Test failed vertex creation is counted in the import result.""" mock_handle_graph_creation.return_value = None vertices = [{"label": "person", "properties": {"name": "Tom Hanks", "age": 67}}] @@ -402,10 +404,30 @@ def test_load_into_graph_raises_explicit_error_when_vertex_creation_fails(self, } ] - with self.assertRaisesRegex(ValueError, "Failed to create vertex"): - self.commit2graph.load_into_graph(vertices, edges, self.schema) + result = self.commit2graph.load_into_graph(vertices, edges, self.schema) - mock_handle_graph_creation.assert_called_once() + self.assertEqual(result["vertices_attempted"], 1) + self.assertEqual(result["vertices_created"], 0) + 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.assertEqual( + result["errors"][0], {"kind": "vertex", "index": 0, "reason": "create_failed", "label": "person"} + ) + self.assertEqual( + result["errors"][1], + { + "kind": "edge", + "index": 0, + "reason": "endpoint_vertex_failed", + "label": "acted_in", + "key": "outV", + }, + ) + self.assertNotIn("Tom Hanks", str(result["errors"])) + self.assertNotIn("Forrest Gump", str(result["errors"])) + self.assertEqual(mock_handle_graph_creation.call_count, 1) @patch("hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph.Commit2Graph._handle_graph_creation") def test_load_into_graph_with_data_type_validation_failure(self, mock_handle_graph_creation): @@ -425,6 +447,14 @@ def test_load_into_graph_with_data_type_validation_failure(self, mock_handle_gra } ] - self.commit2graph.load_into_graph(vertices, edges, self.schema) + result = self.commit2graph.load_into_graph(vertices, edges, self.schema) - self.assertEqual(mock_handle_graph_creation.call_count, 1) + self.assertEqual(result["vertices_created"], 0) + self.assertEqual(result["vertices_skipped"], 2) + self.assertEqual(result["edges_created"], 0) + self.assertEqual(result["edges_skipped"], 1) + self.assertEqual( + result["errors"][-1], + {"kind": "edge", "index": 0, "reason": "endpoint_vertex_failed", "label": "acted_in", "key": "outV"}, + ) + mock_handle_graph_creation.assert_not_called() 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 7ccf7310e..5fd3aeeef 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 @@ -119,7 +119,7 @@ def test_init_uses_request_local_connection_settings(self, mock_settings, mock_c @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): + def test_init_request_config_falls_back_for_none_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" @@ -136,7 +136,29 @@ def test_init_request_config_does_not_inherit_global_graphspace(self, mock_setti ) _, kwargs = mock_client_class.call_args - assert kwargs["graphspace"] is None + assert kwargs["graphspace"] == "global_space" + 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_request_config_preserves_explicit_empty_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": "", + }, + ) + + _, kwargs = mock_client_class.call_args + assert kwargs["graphspace"] == "" assert kwargs["url"] == "10.0.0.1:8080" @patch("hugegraph_llm.operators.hugegraph_op.schema_manager.PyHugeClient") diff --git a/hugegraph-llm/src/tests/operators/llm_op/test_info_extract.py b/hugegraph-llm/src/tests/operators/llm_op/test_info_extract.py index e352a097b..7c8fc189a 100644 --- a/hugegraph-llm/src/tests/operators/llm_op/test_info_extract.py +++ b/hugegraph-llm/src/tests/operators/llm_op/test_info_extract.py @@ -120,6 +120,39 @@ def test_extract_by_regex_with_schema(self): self.assertEqual(actual_edges, expected_edges) self.assertEqual(graph["schema"], self.schema) + def test_extract_by_regex_with_check_schema_shape(self): + schema = { + "vertexlabels": [ + {"name": "person", "properties": ["name", "age", "occupation"]}, + {"name": "webpage", "properties": ["name", "url"]}, + ], + "edgelabels": [ + { + "name": "roommate", + "source_label": "person", + "target_label": "person", + "properties": [], + } + ], + } + graph = {"triples": [], "vertices": [], "edges": [], "schema": schema} + + extract_triples_by_regex_with_schema(schema, self.llm_output, graph) + + self.assertEqual( + graph["edges"], + [{"start": "person-Alice", "end": "person-Bob", "type": "roommate", "properties": {}}], + ) + self.assertIn( + { + "id": "person-Alice", + "name": "Alice", + "label": "person", + "properties": {"name": "Alice", "age": "25", "occupation": "lawyer"}, + }, + graph["vertices"], + ) + def test_extract_by_regex(self): graph = {"triples": []} extract_triples_by_regex(self.llm_output, graph) diff --git a/hugegraph-llm/src/tests/operators/llm_op/test_property_graph_extract.py b/hugegraph-llm/src/tests/operators/llm_op/test_property_graph_extract.py index efb1e28fe..5e4365644 100644 --- a/hugegraph-llm/src/tests/operators/llm_op/test_property_graph_extract.py +++ b/hugegraph-llm/src/tests/operators/llm_op/test_property_graph_extract.py @@ -18,6 +18,7 @@ # pylint: disable=protected-access import json +import time import unittest from unittest.mock import MagicMock, patch @@ -200,6 +201,20 @@ def test_filter_item(self): self.assertEqual(filtered_items[2]["properties"]["role"], "Forrest Gump") self.assertNotIn("ignored", filtered_items[2]["properties"]) + def test_filter_item_accepts_parent_edge_label_without_properties(self): + """Test parent edge labels without properties do not break filtering.""" + schema = { + "vertexlabels": self.schema["vertexlabels"], + "edgelabels": [ + {"name": "BELONGS_TO", "edgelabel_type": "PARENT"}, + *self.schema["edgelabels"], + ], + } + + filtered_items = filter_item(schema, []) + + self.assertEqual(filtered_items, []) + def test_extract_property_graph_by_llm(self): """Test the extract_property_graph_by_llm method.""" extractor = PropertyGraphExtract(llm=self.mock_llm) @@ -1053,6 +1068,46 @@ def test_extract_and_filter_label_missing_keys(self): self.assertEqual(result, []) + def test_extract_and_filter_label_drops_items_with_non_object_properties(self): + """Drop malformed LLM items whose properties field is not an object.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + text = """{ + "vertices": [ + { + "label": "person", + "properties": null + }, + { + "label": "movie", + "properties": { + "title": "Forrest Gump" + } + } + ], + "edges": [ + { + "label": "acted_in", + "properties": ["role"], + "source": { + "label": "person", + "properties": ["Tom Hanks"] + }, + "target": { + "label": "movie", + "properties": { + "title": "Forrest Gump" + } + } + } + ] + }""" + + result = extractor._extract_and_filter_label(self.schema, text) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["type"], "vertex") + self.assertEqual(result[0]["label"], "movie") + def test_run(self): """Test the run method.""" extractor = PropertyGraphExtract(llm=self.mock_llm) @@ -1081,6 +1136,123 @@ def test_run(self): # Check edge properties self.assertEqual(result["edges"][0]["properties"]["role"], "Forrest Gump") + def test_run_with_parallel_chunks_preserves_input_order(self): + """Test parallel chunk extraction keeps deterministic aggregation order.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + + def fake_extract(_schema, chunk): + if chunk == "slow chunk": + time.sleep(0.02) + name = "Alice" if chunk == "slow chunk" else "Bob" + return json.dumps( + { + "vertices": [ + { + "type": "vertex", + "label": "person", + "properties": {"name": name}, + } + ], + "edges": [], + } + ) + + extractor.extract_property_graph_by_llm = MagicMock(side_effect=fake_extract) + context = { + "schema": self.schema, + "chunks": ["slow chunk", "fast chunk"], + "max_parallel_chunks": 8, + } + + result = extractor.run(context) + + self.assertEqual([vertex["properties"]["name"] for vertex in result["vertices"]], ["Alice", "Bob"]) + self.assertEqual(result["call_count"], 2) + self.assertEqual(result["max_parallel_chunks"], 2) + self.assertEqual(extractor.extract_property_graph_by_llm.call_count, 2) + + @patch("hugegraph_llm.operators.llm_op.property_graph_extract.ThreadPoolExecutor") + def test_run_with_single_parallel_chunk_uses_serial_path(self, mock_executor): + """Test max_parallel_chunks=1 keeps the existing serial extraction path.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + extractor.extract_property_graph_by_llm = MagicMock(side_effect=self.llm_responses) + context = { + "schema": self.schema, + "chunks": self.chunks, + "max_parallel_chunks": 1, + } + + result = extractor.run(context) + + mock_executor.assert_not_called() + self.assertEqual(extractor.extract_property_graph_by_llm.call_count, 2) + self.assertEqual(result["call_count"], 2) + + @patch("hugegraph_llm.operators.llm_op.property_graph_extract.ThreadPoolExecutor") + def test_run_with_one_chunk_reduces_parallelism_to_direct_extract(self, mock_executor): + """Test one pre-split chunk is extracted directly even when parallelism is higher.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + extractor.extract_property_graph_by_llm = MagicMock(return_value=self.llm_responses[0]) + context = { + "schema": self.schema, + "chunks": [self.chunks[0]], + "max_parallel_chunks": 4, + } + + result = extractor.run(context) + + mock_executor.assert_not_called() + self.assertEqual(extractor.extract_property_graph_by_llm.call_count, 1) + self.assertEqual(result["call_count"], 1) + self.assertEqual(result["max_parallel_chunks"], 1) + + @patch("hugegraph_llm.operators.llm_op.property_graph_extract.ThreadPoolExecutor") + def test_run_with_empty_chunks_keeps_parallelism_metadata_positive(self, mock_executor): + """Test empty chunks do not record max_parallel_chunks as zero.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + extractor.extract_property_graph_by_llm = MagicMock() + context = { + "schema": self.schema, + "chunks": [], + "max_parallel_chunks": 4, + } + + result = extractor.run(context) + + mock_executor.assert_not_called() + extractor.extract_property_graph_by_llm.assert_not_called() + self.assertEqual(result["call_count"], 0) + self.assertEqual(result["max_parallel_chunks"], 1) + + def test_run_raises_when_chunk_output_is_malformed_json(self): + """Test flow execution fails instead of silently dropping malformed chunk output.""" + extractor = PropertyGraphExtract(llm=self.mock_llm) + extractor.extract_property_graph_by_llm = MagicMock(return_value='{"vertices": [], "edges": [}') + context = { + "schema": self.schema, + "chunks": ["malformed output chunk"], + } + + with self.assertRaisesRegex(ValueError, "Invalid property graph JSON") as error: + extractor.run(context) + + self.assertIsInstance(error.exception.__cause__, json.JSONDecodeError) + + def test_run_falls_back_to_default_parallelism_for_invalid_context_value(self): + """Test invalid operator-level parallelism input does not fail before extraction.""" + extractor = PropertyGraphExtract(llm=self.mock_llm, max_parallel_chunks=2) + extractor.extract_property_graph_by_llm = MagicMock(return_value=self.llm_responses[0]) + context = { + "schema": self.schema, + "chunks": [self.chunks[0]], + "max_parallel_chunks": "not-an-int", + } + + result = extractor.run(context) + + self.assertEqual(result["call_count"], 1) + self.assertEqual(result["max_parallel_chunks"], 1) + def test_run_with_existing_vertices_and_edges(self): """Test the run method with existing vertices and edges.""" extractor = PropertyGraphExtract(llm=self.mock_llm) @@ -1123,5 +1295,95 @@ def test_run_with_existing_vertices_and_edges(self): self.assertEqual(result["edges"][0]["properties"]["role"], "Jack Dawson") +def test_run_debug_logs_do_not_include_raw_chunk_or_llm_output(): + schema = { + "vertexlabels": [ + { + "id": 1, + "name": "person", + "primary_keys": ["name"], + "nullable_keys": [], + "properties": ["name"], + } + ], + "edgelabels": [], + } + sensitive_chunk = "Alice secret source document" + sensitive_output = json.dumps( + { + "vertices": [ + { + "type": "vertex", + "label": "person", + "properties": {"name": "Alice secret model output"}, + } + ], + "edges": [], + } + ) + extractor = PropertyGraphExtract(llm=MagicMock(spec=BaseLLM)) + extractor.extract_property_graph_by_llm = MagicMock(return_value=sensitive_output) + + with patch("hugegraph_llm.operators.llm_op.property_graph_extract.log.debug") as debug_log: + extractor.run({"schema": schema, "chunks": [sensitive_chunk]}) + + logged_args = " ".join(str(arg) for call in debug_log.call_args_list for arg in call.args) + assert "chunk processed" in logged_args + assert sensitive_chunk not in logged_args + assert "Alice secret model output" not in logged_args + + +def test_invalid_edge_warning_does_not_include_raw_edge_payload(): + schema = { + "vertexlabels": [ + { + "id": 1, + "name": "person", + "primary_keys": ["name"], + "nullable_keys": [], + "properties": ["name"], + }, + { + "id": 2, + "name": "movie", + "primary_keys": ["title"], + "nullable_keys": [], + "properties": ["title"], + }, + ], + "edgelabels": [{"name": "acted_in", "properties": ["role"], "source_label": "person", "target_label": "movie"}], + } + graph_json = json.dumps( + { + "vertices": [ + { + "type": "vertex", + "label": "person", + "properties": {"name": "Alice secret vertex"}, + } + ], + "edges": [ + { + "type": "edge", + "label": "acted_in", + "properties": {"role": "sensitive role"}, + "source": {"label": "person", "properties": {"name": "Alice secret vertex"}}, + "target": {"label": "movie", "properties": {"title": "Missing secret movie"}}, + } + ], + } + ) + extractor = PropertyGraphExtract(llm=MagicMock(spec=BaseLLM)) + + with patch("hugegraph_llm.operators.llm_op.property_graph_extract.log.warning") as warning_log: + extractor._extract_and_filter_label(schema, graph_json, raise_on_invalid=True) + + logged_args = " ".join(str(arg) for call in warning_log.call_args_list for arg in call.args) + assert "Invalid edge endpoints" in logged_args + assert "sensitive role" not in logged_args + assert "Alice secret vertex" not in logged_args + assert "Missing secret movie" not in logged_args + + if __name__ == "__main__": unittest.main() diff --git a/hugegraph-llm/src/tests/utils/test_graph_index_utils.py b/hugegraph-llm/src/tests/utils/test_graph_index_utils.py new file mode 100644 index 000000000..0a8f2ec3e --- /dev/null +++ b/hugegraph-llm/src/tests/utils/test_graph_index_utils.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. + +from unittest.mock import Mock + +from hugegraph_llm.flows import FlowName +from hugegraph_llm.utils import graph_index_utils + + +def test_extract_graph_gradio_path_keeps_existing_scheduler_call(monkeypatch): + scheduler = Mock() + scheduler.schedule_flow.return_value = '{"vertices":[],"edges":[]}' + monkeypatch.setattr(graph_index_utils, "read_documents", Mock(return_value=["marko knows vadas"])) + monkeypatch.setattr(graph_index_utils.SchedulerSingleton, "get_instance", Mock(return_value=scheduler)) + + result = graph_index_utils.extract_graph(None, "marko knows vadas", "hugegraph", "prompt") + + assert result == '{"vertices":[],"edges":[]}' + scheduler.schedule_flow.assert_called_once_with( + FlowName.GRAPH_EXTRACT, + "hugegraph", + ["marko knows vadas"], + "prompt", + "property_graph", + split_type="document", + )