feat(llm): expand graph extraction service APIs

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

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

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

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,
         )

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

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

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):

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:

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)